analyze_mutationtree.py

import sys
import os
import copy
import traceback
import simplejson as json
import random
import math
import matplotlib.pyplot as plt
from matplotlib.collections import PatchCollection
from matplotlib.patches import Rectangle
import numpy as np
import pygraphviz as PG
# from igraph import Graph, EdgeSeq

from plot import plot, plotBar, plotCDF, plotBar1, plotCDF2
from utils import loadDataCached, getTestParams
from analyze_programs import Program, __processTest 

class Node:
    def __init__(self, p: Program):
        self.program = p
        self.level = -1
        self.height = -1 # Starting from this node, the max height
        self.size = 1
        self.children = []
        self.parent = None
        self.id = None
    def computeLevel(self, level=-1, visited=None, maxlevel=10000000):
        if level >= 0:
            self.level = level
        height = self.level
        self.size = 1
        if self.level >= maxlevel:
            self.children = []
        for n in self.children:
            if visited is not None and n.program.sig in visited:
                print("Cycle detected:", n.program.sig)
                continue
            n.computeLevel(self.level + 1, visited, maxlevel)
            d = n.height
            s = n.size
            self.size += s
            height = d if d > height else height
        self.height = height
        return height
    def collectDegrees(self, result, visited=None):
        degree = 0
        leaves = 0
        for n in self.children:
            if visited is not None and n.program.sig in visited:
                print("Cycle detected:", n.program.sig)
                continue
            degree += 1
            l = n.collectDegrees(result, visited=visited)
            leaves += l
        if degree > 0: # Ignore leaf nodes
            result.append(degree)
            return leaves
        else:
            return 1
    def getOffspring(self):
        ret = [self]
        for n in self.children:
            ret += n.getOffspring()
        return ret
        
    #def __eq__(self, other):
    #    return self.p.sig == self.other.sig

def buildMutationTrees(p_all, p_generated, p_corpus, p_triage):
    roots = []
    nodes_all_dedup = {}
    # Establish roots
    for pid in p_generated:
        p = p_all[pid]
        if not p.sig in nodes_all_dedup:
            n = Node(p)
            n.level = 0
            nodes_all_dedup[p.sig] = n 
            roots.append(n)
    # Back trace to build the tree
    for i in range(len(p_all)):
        p = p_all[i]
        if not p.executed:
            continue
        if p.sig in nodes_all_dedup:
            continue
        # Ignore minimize nodes
        if p.origin == "Minimize":
            continue
        # Backtrack
        p_cur = p
        n_cur = Node(p)
        nodes_all_dedup[p.sig] = n_cur
        n_parent = None
        p_parent = None
        root_reached = True
        while p_cur.parent is not None:
            p_parent = p_all[p_cur.parent]
            dist = 1
            while p_parent.origin == "Minimize" and p_parent.parent is not None: # Ignore minimize edges
                p_parent = p_all[p_parent.parent]
                dist += 1
            #print(p_parent.sig, p_cur.sig, dist)
            #exit()
            if p_parent.sig in nodes_all_dedup: # This check ensures there will be no cycles
                n_parent = nodes_all_dedup[p_parent.sig]
                n_cur.parent = n_parent
                n_parent.children.append(n_cur)
                root_reached = False
                break
            else:
                n_parent = Node(p_parent)
                nodes_all_dedup[p_parent.sig] = n_parent
                n_cur.parent = n_parent
                n_parent.children.append(n_cur)
                p_cur = p_parent
                n_cur = n_parent
    # Sort out degree from root
    for n in roots:
        nodes_visited = set() # Check whether there are cycles
        n.computeLevel(visited=nodes_visited)
        # print(n.program.sig, n.height, len(n.children), len(n.program.children), len(n.program.minimizeChildren))
    __roots = []
    # Remove isolated roots
    for n in roots:
        if len(n.children) > 0:
            __roots.append(n)
    return __roots, nodes_all_dedup

def sortSample(l, n):
    if n == 0:
        return []
    tmp = sorted(l, reverse=True, key=lambda x: x.size)
    step = len(l) / float(n)
    ret = []
    for i in range(n):
        idx = round(step * (i))
        if 0 <= idx < len(tmp):
            ret.append(tmp[idx])
        else:
            ret.append(tmp[-1])
    return ret

def sampleTreeNoAdd(trees, nodes, samplerate=0.05):
    # Node sampling
    nodes_sample = random.sample(nodes.keys(), math.ceil(len(nodes) * samplerate))
    # Reconstruct sampled tree via back-tracking 
    nodes_sample_dedup = {}
    # Add all nodes to nodes
    for n in nodes_sample:
        nodes_sample_dedup[nodes[n].program.sig] = None
    trees_sample = []
    for nsig in nodes_sample:
        #if nsig in nodes_sample_dedup:
        #    continue
        n_cur = nodes[nsig]
        n_new_cur = Node(nodes[nsig].program)
        nodes_sample_dedup[nsig] = n_new_cur
        n_parent = n_cur.parent
        root_reached = True
        while n_parent is not None:
            n_parent_sig = n_parent.program.sig
            if n_parent_sig in nodes_sample_dedup: # This check ensures there will be no cycles
                print(n_parent_sig,'->',n_cur.program.sig)
                if nodes_sample_dedup[n_parent_sig] is None:
                    nodes_sample_dedup[n_parent_sig] = Node(n_parent.parent)
                    n_new_parent = nodes_sample_dedup[n_parent_sig]
                    n_new_cur.parent = n_new_parent
                    n_new_parent.children.append(n_new_cur)
                else:
                    n_new_parent = nodes_sample_dedup[n_parent_sig]
                    n_new_cur.parent = n_new_parent
                    n_new_parent.children.append(n_new_cur)
                    root_reached = False
                    break
            else:
                n_cur = n_parent
                n_parent = n_cur.parent
        print(root_reached)
        if root_reached:
            #if len(n_new_cur.children) > 0:
                trees_sample.append(n_new_cur)
    for n in trees_sample:
        nodes_visited = set() # Check whether there are cycles
        n.computeLevel(visited=nodes_visited)
    print(trees_sample, len(nodes_sample_dedup))
    return trees_sample, nodes_sample_dedup


def sampleTree(trees, nodes, samplerate=0.05, maxlevel=100):
    # Node sampling
    nodes_sample = random.sample(nodes.keys(), math.ceil(len(nodes) * samplerate))
    # Reconstruct sampled tree via back-tracking 
    nodes_sample_dedup = {}
    trees_sample = []
    for nsig in nodes_sample:
        if nsig in nodes_sample_dedup:
            continue
        n_cur = nodes[nsig]
        n_new_cur = Node(nodes[nsig].program)
        nodes_sample_dedup[nsig] = n_new_cur
        n_parent = n_cur.parent
        root_reached = True
        while n_parent is not None:
            n_parent_sig = n_parent.program.sig
            if n_parent_sig in nodes_sample_dedup: # This check ensures there will be no cycles
                n_new_parent = nodes_sample_dedup[n_parent_sig]
                n_new_cur.parent = n_new_parent
                n_new_parent.children.append(n_new_cur)
                root_reached = False
                break
            else:
                n_new_parent = Node(n_parent.program)
                nodes_sample_dedup[n_parent_sig] = n_new_parent
                n_new_cur.parent = n_new_parent
                n_new_parent.children.append(n_new_cur)
                n_new_cur = n_new_parent
                n_cur = n_parent
                n_parent = n_cur.parent
        if root_reached and n_new_cur.program.origin == "Generate":
            if len(n_new_cur.children) > 0:
                trees_sample.append(n_new_cur)
    for n in trees_sample:
        nodes_visited = set() # Check whether there are cycles
        n.computeLevel(visited=nodes_visited, maxlevel=maxlevel)
    print(trees_sample, len(nodes_sample_dedup))
    return trees_sample, nodes_sample_dedup

def plotForest(trees, nodes, samplerate=1.0, prunerate=0.0, maxnodes=500, maxtrees=100000, maxlevel=5000, outfile="tree.png"):
    #if maxnodes / len(nodes) < samplerate / 2:
    samplerate = maxnodes * 2.0 / len(nodes)
    AG = PG.AGraph(directed=True, strict=True)
    AG.graph_attr.update(penwidth=3, ratio='fill', size='10,16')
    AG.node_attr.update(color='black', shape='circle', style='filled', fillcolor='red',penwidth=3)
    AG.edge_attr.update(len='2.0',penwidth=3)
    if samplerate < 1.0:
        trees_sample, nodes_sample_dedup = sampleTree(trees, nodes, samplerate=samplerate, maxlevel=maxlevel)
    else:
        trees_sample = trees
    # Prune randomly if necessary
    size_total = 0
    for r in trees_sample:
        size_total += r.size
    print(size_total)
    while size_total > maxnodes:
        nsig = random.choice(list(nodes_sample_dedup.keys()))
        n = nodes_sample_dedup[nsig]
        # Remove 
        if n.parent is not None and n in n.parent.children:
            n.parent.children.remove(n)
        if n in trees_sample:
            trees_sample.remove(n)
        # Recompute sizes and everything
        size_total = 0
        for r in trees_sample:
            nodes_visited = set() # Check whether there are cycles
            r.computeLevel(visited=nodes_visited, maxlevel=maxlevel)
            size_total += r.size
        # Remove nodes from pool
        deleted_nodes = n.getOffspring()
        print("Deleting:", n, n.size)
        for dn in deleted_nodes:
            dnsig = dn.program.sig
            if dnsig in nodes_sample_dedup:
                del nodes_sample_dedup[dnsig]
        #for n in trees_sample:
        #    nodes_visited = set() # Check whether there are cycles
        #    n.computeLevel(visited=nodes_visited, maxlevel=maxlevel)
    print(size_total)
    '''
    if len(trees_sample) > 0:
        n_trees = min(math.ceil(len(trees_sample) * (1.0 - prunerate)), maxtrees)
        # n_trees = 1 + round(math.log(len(trees_sample)))
        n_trees = math.ceil(math.pow(len(trees_sample), 0.33))
        trees_sample = sortSample(trees_sample, n_trees)
    '''
    # Plot trees, prune along the way
    queue = []
    nid = 0
    for root in trees_sample:
        if len(root.children) == 0: # Ignore stray nodes
            continue
        root.level = 1
        queue.append(root)
        root.id = nid
        nid += 1
    cur_level = 1
    while len(queue) > 0 and maxnodes > 0:
        # print([x.id for x in queue])
        node = queue.pop(0)
        # Check if we need pruning
        #if len(queue) > 0 and queue[0].level > cur_level:
        #    n_nodes = math.ceil(len(queue) * (1.0 - prunerate))
        #    queue = sortSample(queue, n_nodes)
        AG.add_node(node.id, label='')
        maxnodes -= 1
        if node.level < maxlevel and len(node.children) > 0:
          n_childs = math.ceil(len(node.children) * (1.0 - prunerate))
          # n_childs = 1 + round(math.log(len(node.children)))
          # n_childs = math.ceil(math.pow(len(node.children), 0.33))
          children_subset = sortSample(node.children, n_childs)
          for c in children_subset:
            queue.append(c)
            c.id = nid
            c.level = node.level + 1
            c.parent = node
            nid += 1
        # Edge to parent
        if node.parent is not None:
            parent_id = node.parent.id
            if parent_id is None:
                print("WTF")
            AG.add_edge(parent_id, node.id)
            print(parent_id, node.id)

    '''
    n_trees = math.ceil(len(trees) * samplerate)
    print(n_trees)
    trees_subset = sortSample(trees, n_trees)
    print(trees_subset)
    # Construct tree in BFS fashion
    queue = []
    nid = 0
    for root in trees_subset:
        queue.append(root)
        root.id = nid
        nid += 1
    while len(queue) > 0 and maxnodes > 0:
        node = queue.pop(0)
        AG.add_node(nid)
        maxnodes -= 1
        # Child sampling
        n_childs = math.ceil(len(node.children) * samplerate)
        children_subset = sortSample(node.children, n_childs)
        # Add children to queue
        for c in children_subset:
            queue.append(c)
            c.id = nid
            nid += 1
        # Edge to parent
        if node.parent is not None:
            parent_id = node.parent.id
            if parent_id is None:
                print("WTF")
            AG.add_edge(parent_id, node.id)
            print(parent_id, node.id)
    '''
    # Draw tree
    AG.layout(prog='dot')
    AG.draw(outfile, format='png', prog='dot')
    AG.draw(outfile + ".pdf", format='pdf', prog='dot')

def plotMutationTree(tests):
    datas = {}
    for test in tests:
        name, module, run = getTestParams(test)
        # name = name + '_' + module
        datas[name] = {
            "Num_Trees": [],
            "Tree_Height": [],
            "Node_Degree": [],
            "Leaf_Nodes": [],
            "Leaf_Nodes_Percentage": [],
            "Leaf_Nodes_Per_Tree": [],
            "Leaf_Nodes_Percentage_Per_Tree": [],
            "Tree_Size": [],
            "Seed_Subtree_Height": {"Generate": [], "Mutate": [], "Minimize": [], "All": []},
            "Seed_Subtree_Size": {"Generate": [], "Mutate": [], "Minimize": [], "All": []}
        }
    for test in tests:
        name, module, run = getTestParams(test)
        # name = name + '_' + module
        print("Plotting mutation tree for %s" % test)
        try:
            p_all, p_generated, p_corpus, p_triage, __data = loadDataCached('program_%s.cache', test, __processTest);
            for i in range(len(p_all)):
                if type(p_all[i]) == dict:
                    p_all[i] = Program.FromDict(p_all[i])
        except:
            traceback.print_exc()
            continue;
        trees, nodes = buildMutationTrees(p_all, p_generated, p_corpus, p_triage)
        print(len(trees), len(nodes))
        plotForest(trees, nodes, outfile="mt_sample_%s.png" % test)
        datas[name]["Num_Trees"].append(len(trees))
        num_lf = 0
        for r in trees:
            # For generation, just consider ones that're mutated
            if r.height > 0 and r.size > 1:
                datas[name]["Tree_Height"].append(r.height)
                datas[name]["Tree_Size"].append(r.size)
            degrees = []
            visited = set()
            l = r.collectDegrees(degrees, visited)
            # print(l, len(degrees))
            datas[name]["Node_Degree"] += degrees
            num_lf += l
            datas[name]["Leaf_Nodes_Per_Tree"].append(l)
            datas[name]["Leaf_Nodes_Percentage_Per_Tree"].append(l / r.size)
        datas[name]["Leaf_Nodes"].append(num_lf)
        if len(nodes) > 0:
            datas[name]["Leaf_Nodes_Percentage"].append(100.0 * num_lf / len(nodes))
        for sig in nodes:
            n = nodes[sig]
            # datas[name]["Node_Degree"].append(len(n.children))
            csig = n.program.sig
            if csig in p_corpus:
                src = p_all[p_corpus[csig]].corpusSource
                if src is not None:
                    datas[name]["Seed_Subtree_Height"][src].append(n.height)
                    datas[name]["Seed_Subtree_Size"][src].append(n.size)
                    datas[name]["Seed_Subtree_Height"]["All"].append(n.height)
                    datas[name]["Seed_Subtree_Size"]["All"].append(n.size)
    # Overall Tree size and height
    # plotCDF(datas, key="Tree_Height", xlabel="Tree Height", ylabel="Cumulative # Trees", title="", outfile="mt_height_overall.png", xlogscale=True, raw=True, small=False);
    # plotCDF(datas, key="Tree_Size", xlabel="Tree Size", ylabel="Cumulative # Trees", title="", outfile="mt_size_overall.png", xlogscale=True, raw=True, small=False);
    plotCDF2([datas, datas], key=["Tree_Size", "Tree_Height"], xlabel=["Tree Size", "Tree Height"], ylabel=["Cumulative # Trees", "Cumulative # Trees"], outfile="mt_size_height_overall.png", xlogscale=[True, True], raw=[True, True]);
    # plotCDF(datas, key="Node_Degree", xlabel="Node Degree", ylabel="Cumulative # Nodes", title="", outfile="mt_degree_overall.png", xlogscale=True, raw=True);

    #plotCDF(datas, key="Tree_Height", xlabel="Tree Height", ylabel="CDF", title="", outfile="mt_height_overall_cdf.png", xlogscale=False, xrange=(-5,105));
    #plotCDF(datas, key="Tree_Size", xlabel="Tree Size", ylabel="CDF", title="", outfile="mt_size_overall_cdf.png", xlogscale=True);
    #plotCDF(datas, key="Node_Degree", xlabel="Tree Size", ylabel="CDF", title="", outfile="mt_degree_overall_cdf.png", xlogscale=True);

    # Leaf nodes per tree
    #plotCDF(datas, key="Leaf_Nodes_Per_Tree", xlabel="# of nodes", ylabel="CDF", title="", outfile="mt_leaf_per_tree_cdf.png", xlogscale=True);
    #plotCDF(datas, key="Leaf_Nodes_Percentage_Per_Tree", xlabel="# of nodes", ylabel="CDF", title="", outfile="mt_leaf_per_tree_perc_cdf.png", xlogscale=False);

    # Num Trees
    tmp = {}
    for name in datas:
        #tmp["Num_Trees"][name] = datas[name]["Num_Trees"]
        tmp[name] = {"Num_Trees": datas[name]["Num_Trees"]}
    plotBar1(tmp, ylabel="# Nodes", outfile="mt_num_trees.png")


    # Leaf nodes total
    tmp = {}
    for name in datas:
        #tmp["Leaf_Nodes"][name] = datas[name]["Leaf_Nodes"]
        tmp[name] = {"Leaf_Nodes": datas[name]["Leaf_Nodes"]}
    plotBar1(tmp, ylabel="# Nodes", outfile="mt_leaf_per_forest.png")
    tmp = {"Leaf_Nodes_Percentage": {}}
    for name in datas:
        #tmp["Leaf_Nodes_Percentage"][name] = datas[name]["Leaf_Nodes_Percentage"]
        tmp[name] = {"Leaf_Nodes_Percentage": datas[name]["Leaf_Nodes_Percentage"]}
    plotBar1(tmp, ylabel="Percentage (%)", outfile="mt_leaf_per_forest_perc.png")

    # Sub-tree size, height
    tmp = {"Generate": {}, "Mutate": {}, "Minimize": {}, "All": {}}
    for src in ["Generate", "Mutate", "Minimize", "All"]:
        for name in datas:
            tmp[src][name] = datas[name]["Seed_Subtree_Height"][src]
        plotCDF(tmp[src], xlabel="Tree height", ylabel="CDF", title="", outfile="mt_subtree_height_cdf_%s.png" % src, xlogscale=False);
    tmp = {"Generate": {}, "Mutate": {}, "Minimize": {}, "All": {}}
    for src in ["Generate", "Mutate", "Minimize", "All"]:
        for name in datas:
            tmp[src][name] = datas[name]["Seed_Subtree_Size"][src]
        plotCDF(tmp[src], xlabel="Tree height", ylabel="CDF", title="", outfile="mt_subtree_size_cdf_%s.png" % src, xlogscale=True);