EGoT_KAnonymity/tree_handler.py at main · PortlandStatePowerLab/EGoT_KAnonymity · 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
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
from treelib import Node, Tree
import numpy as np
import pandas as pd


class Dist:
    def __init__(self,root_name='substation',max_size=100000):
        # generate new tree
        self.tree = Tree()
        self.id = 0 # will be increment counter that used for ids
        # create root feeder
        root_id = f'{root_name}{self.id}'
        self.tree.create_node(root_name,root_id)
        self.current_lvl = [root_id]
        self.max_cap = max_size
        self.current_cap = 0
        self.lvl_names = [root_name]
        self.leaves = 0
        pass
    def __repr__(self):
        # prints the tree
        self.tree.show()
        return '-'*50
    def add_level(self,lvl_name,k=1,leave=False,lvls=None,id_range=None):
        '''
            Add a level to the current level of the tree.
            It randomly picks which node to
            lvl_name: name used to prefix the level
            k: number of instances to generate
        '''
        lvl = []
        current_lvl = self.current_lvl.copy() if lvls==None else lvls
        # print('current level: ',current_lvl)
        r = list(range(k)) if id_range == None else id_range
        if not lvl_name in self.lvl_names:
            self.lvl_names.append(lvl_name)
        for parent in current_lvl:
            for i in r:
                if leave:
                    name = f'{parent}-{lvl_name}{self.leaves}'
                    self.leaves += 1
                else:
                    name = f'{parent}-{lvl_name}{i}' #if not leave else f'{parent}-{lvl_name}{self.leave}'
                # randomly pick which node to add the new node to
                # parent = np.random.choice(current_lvl)
                # parent = current_lvl
                self.tree.create_node(name,name,parent)
                lvl.append(name)
                self.current_cap += 1
                assert self.current_cap <= self.max_cap, 'Exceeding max number of nodes threshold'
        # use the newly added lvl as the current lvl
        # old_lvl = self.current_lvl
        self.current_lvl = lvl
        return (True,lvl)
    def add_level_rand(self,lvl_name,k=1,leave=False,id_range=None):
        '''
            Add a level to the current level of the tree.
            It randomly picks which node to
            lvl_name: name used to prefix the level
            k: number of instances to generate
        '''
        lvl = []
        current_lvl = self.current_lvl.copy()
        # print('current level: ',current_lvl)
        if not lvl_name in self.lvl_names:
            self.lvl_names.append(lvl_name)
        for i in range(k):
            # randomly pick which node to add the new node to
            parent = np.random.choice(current_lvl)
            if len(current_lvl) > 1:
                current_lvl.remove(parent)
            if leave:
                name = f'{parent}-{lvl_name}{self.leaves}'
                self.leaves += 1
            else:
                name = f'{parent}-{lvl_name}{i}' #if not leave else f'{parent}-{lvl_name}{self.leave}'
            self.tree.create_node(name,name,parent)
            lvl.append(name)
            self.current_cap += 1
            assert self.current_cap <= self.max_cap, 'Exceeding max number of nodes threshold'
        self.current_lvl = lvl
        return (True,lvl)
    def drop_childless(self):
        '''
            Cleans the tree from unused paths (balances the tree so that all leaves on the same level).
            It also renumbers the remaining levels.
        '''
        # self.tree = self.tree.filter_nodes(lambda x: x.is)
        pass
    def export_to_df(self):
        '''
            Performs depth first traversal to get the a representation of the tree.
            It exports the rep. tree into a df with given headers.
            * Headers: desired header names (must be equal to the tree depth)
            NOTE: It drops anything that isn't on the same level (i.e. only the balanced parts )
        '''
        # assert len(headers) == self.tree.depth, 'Incompatible '
        depth = self.tree.depth() + 1 # including the root
        tree = self.tree.paths_to_leaves()
        tree = list(filter(lambda x: len(x) == depth,tree))
        df = pd.DataFrame(tree,columns = self.lvl_names)
        return df
    def drop_subtrees(self,trees):
        '''
            drops the subtrees given in the args
        '''
        for i in trees:
            for c in self.tree.get_node(i).children():
                self.tree.remove_subtree()
        return True