from abc import ABC from graphviz import Digraph import re import random from termcolor import cprint class Node: def __init__(self, id, label, type, parent=None): """ type: type must be "param" or "value" for type "param" one of the children must be chosen for type "value" all of the children must be set """ self.id = id self.label = label self.parent = parent self.children = [] self.type = type def __repr__(self): return f"{self.id}({self.type})-'{self.label}'" class ParameterTree(ABC): def __init__(self): self.last_node_id = 0 self.create_digraph() self.nodes = {} self.root = None def create_digraph(self): self.tree = Digraph("unix", format='svg') self.tree.attr(size='30,100') def add_node(self, label, parent=None, type="param"): """ All children of this node must be set """ if type not in ["param", "value"]: raise ValueError('Node type must be "param" or "value"') if parent is None and self.root is not None: raise ValueError("Root already set: {}. parent cannot be None. ".format(self.root.id)) # add to graph node_id = self.new_node_id() self.nodes[node_id] = Node(id=node_id, label=label, parent=parent, type=type) if parent is None: self.root = self.nodes[node_id] else: self.nodes[parent.id].children.append(self.nodes[node_id]) return self.nodes[node_id] def sample_env_params(self, ACL=None): parameters = {} nodes = [self.root] # BFS while nodes: node = nodes.pop(0) if node.type == "param": if len(node.children) == 0: raise ValueError("Node {} doesn't have any children.".format(node.label)) if ACL is None: # choose randomly chosen = random.choice(node.children) else: # let the ACL choose chosen = ACL.choose(node, parameters) assert chosen.type == "value" nodes.append(chosen) parameters[node] = chosen elif node.type == "value": nodes.extend(node.children) else: raise ValueError('Node type must be "param" or "value" and is {}'.format(node.type)) return parameters def new_node_id(self): new_id = self.last_node_id + 1 self.last_node_id = new_id return str("node_"+str(new_id)) def print_tree(self, selected_parameters={}): print("Parameter tree") nodes = [self.root] color = None # BFS while nodes: node = nodes.pop(0) if node.type == "param": if node in selected_parameters.keys(): color = "blue" else: color = None if node.parent is not None: cprint("{}: {} ({}) -----> {}: {} ({})".format( node.parent.type, node.parent.label, node.parent.id, node.type, node.label, node.id ), color) else: cprint("{}: {} ({})".format(node.type, node.label, node.id), color) nodes.extend(node.children) def draw_tree(self, filename, selected_parameters={}, ignore_labels=[], folded_nodes=[], label_parser={}, save=True): self.create_digraph() nodes = [self.root] color_param = "grey60" color_value = "lightgray" fontcolor = "black" fontsize = "18" dots_fontsize = "30" folded_param = "grey95" folded_value = "grey95" folded_fontcolor = "gray70" def add_fold_symbol(label, folded=False): return label # return label + " ❯" if folded else label # BFS - construct vizgraph while nodes: node = nodes.pop(0) while node.label in ignore_labels: node = nodes.pop(0) if node.label in folded_nodes: n_label = label_parser.get(node.label, node.label) n_label = add_fold_symbol(n_label, folded=True) if node.type == "param": color = folded_param self.tree.attr('node', shape='box', style="filled", color=color, fontcolor=folded_fontcolor, fontsize=fontsize) self.tree.node(name=node.id, label=n_label, type="parameter") elif node.type == "value": color = folded_value self.tree.attr('node', shape='ellipse', style='filled', color=color, fontcolor=folded_fontcolor, fontsize=fontsize) self.tree.node(name=node.id, label=n_label, type="value") else: raise ValueError(f"Undefined node type {node.type}") # add folded node sign folded_node_id = node.id+"_fold" # self.tree.attr('node', shape='ellipse', style='filled', color="white", fontcolor=folded_fontcolor, fontsize=fontsize) # self.tree.attr('node', shape='none', style='filled', color="gray", fontcolor=folded_fontcolor, fontsize=dots_fontsize) self.tree.attr('node', shape='none', color="white",fontcolor=folded_fontcolor, fontsize=dots_fontsize) self.tree.node(name=folded_node_id, label="...", type="value") self.tree.edge(node.id, folded_node_id, color=folded_fontcolor) elif node.type == "param": if node.label in selected_parameters.keys() and (node == self.root or node.parent.selected): color = "lightblue3" node.selected=True else: color = color_param node.selected=False n_label = label_parser.get(node.label, node.label) n_label = add_fold_symbol(n_label, folded=False) self.tree.attr('node', shape='box', style="filled", color=color, fontcolor=fontcolor, fontsize=fontsize) self.tree.node(name=node.id, label=n_label, type="parameter") nodes.extend(node.children) elif node.type == "value": if (selected_parameters.get(node.parent.label, "Not existent") == node.label) and (node == self.root or node.parent.selected): # if node.label in selected_parameters.values() and (node == self.root or node.parent.selected): color = "lightblue2" node.selected = True else: color = color_value node.selected = False n_label = label_parser.get(node.label, node.label) n_label = add_fold_symbol(n_label, folded=False) # add to vizgraph self.tree.attr('node', shape='ellipse', style='filled', color=color, fontcolor=fontcolor, fontsize=fontsize) self.tree.node(name=node.id, label=n_label, type="value") nodes.extend(node.children) else: raise ValueError(f"Undefined node type {node.type}") if node.parent is not None: self.tree.edge(node.parent.id, node.id) # draw image if save: self.tree.render(filename) print("Tree image saved in : {}".format(filename)) if __name__ == '__main__': # demo of how to use the ParameterTree class tree = ParameterTree() env_type_nd = tree.add_node("Env_type", type="param") inf_seeking_nd = tree.add_node("Information_seeking", parent=env_type_nd, type="value") collab_nd = tree.add_node("Collaboration", parent=env_type_nd, type="value") perc_inf_nd = tree.add_node("Perception_inference", parent=env_type_nd, type="value") raise DeprecationWarning("deprecated parameters") # Information seeking scaffolding_nd = tree.add_node("Scaffolding", parent=inf_seeking_nd, type="param") tree.add_node("lot", parent=scaffolding_nd, type="value") tree.add_node("medium", parent=scaffolding_nd, type="value") tree.add_node("little", parent=scaffolding_nd, type="value") tree.add_node("none", parent=scaffolding_nd, type="value") prag_fr_compl_nd = tree.add_node("Pragmatic_frame_complexity", parent=inf_seeking_nd, type="param") tree.add_node("Eye contact", parent=prag_fr_compl_nd, type="value") tree.add_node("Hello", parent=prag_fr_compl_nd, type="value") emulation_nd = tree.add_node("Emulation", parent=inf_seeking_nd, type="param") tree.add_node("N", parent=emulation_nd, type="value") tree.add_node("Y", parent=emulation_nd, type="value") pointing_nd = tree.add_node("Pointing", parent=inf_seeking_nd, type="param") tree.add_node("No", parent=pointing_nd, type="value") tree.add_node("Direct", parent=pointing_nd, type="value") tree.add_node("Indirect", parent=pointing_nd, type="value") language_graounding_nd = tree.add_node("Language_grounding", parent=inf_seeking_nd, type="param") tree.add_node("No", parent=language_graounding_nd, type="value") tree.add_node("Color", parent=language_graounding_nd, type="value") tree.add_node("Feedback", parent=language_graounding_nd, type="value") problem_nd = tree.add_node("Problem", parent=inf_seeking_nd, type="param") tree.add_node("Boxes", parent=problem_nd, type="value") tree.add_node("Switches", parent=problem_nd, type="value") tree.add_node("Corridors", parent=problem_nd, type="value") obstacles_nd = tree.add_node("Obstacles", parent=inf_seeking_nd, type="param") tree.add_node("no", parent=obstacles_nd, type="value") tree.add_node("lava", parent=obstacles_nd, type="value") tree.add_node("walls", parent=obstacles_nd, type="value") # Collaboration colab_type_nd = tree.add_node("Collaboration type", parent=collab_nd, type="param") tree.add_node("Door Lever", parent=colab_type_nd, type="value") tree.add_node("Door Button", parent=colab_type_nd, type="value") tree.add_node("Marble Run", parent=colab_type_nd, type="value") tree.add_node("Marble Pass", parent=colab_type_nd, type="value") role_nd = tree.add_node("Role", parent=collab_nd, type="param") tree.add_node("A", parent=role_nd, type="value") tree.add_node("B", parent=role_nd, type="value") tree.add_node("asocial", parent=role_nd, type="value") # Perception inference NPC_movement_nd = tree.add_node("NPC movement", parent=perc_inf_nd, type="param") tree.add_node("can't turn; can't move", parent=NPC_movement_nd, type="value") tree.add_node("can turn; can't move", parent=NPC_movement_nd, type="value") tree.add_node("can turn; can move", parent=NPC_movement_nd, type="value") occlusion_nd = tree.add_node("Occlusions", parent=perc_inf_nd, type="param") tree.add_node("no", parent=occlusion_nd, type="value") tree.add_node("walls", parent=occlusion_nd, type="value") params = tree.sample_env_params() tree.draw_tree("viz/demotree", params)