box-sim/utils.py at master · TDoGoodT/box-sim · 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
111
112
113
114
115
116
117
118
119
120
121
import networkx as nx
import matplotlib.pyplot as plt
import numpy as np
from scipy.spatial.transform import Rotation as R
from option_tree import OptionsTree

def_arr = [
    0,
    1,
    1,
    1,
    1,
    0,
    1,
    0,
    1,
    1,
    1,
    1,
    1,
    1,
    1,
    1,
    1,
    1,
    1,
    0,
    1,
    0,
    1,
    1,
    1,
    1,
    0,
]
def get_full_states_tree_df(depth):
    G = OptionsTree()
    options = ['0']
    new_options = []

    for i in range(1,depth-1):
        for option in options:
            for j in range(4):
                new_option = option + str(j)
                new_options.append(new_option)
                G.addEdge(option, new_option)
                # print(option +","+ new_option)

        options = new_options
        new_options = []
    for option in options:
            new_option = option + str(0)
            new_options.append(new_option)
            G.addEdge(option, new_option)
            # print(option +","+ new_option)


    return G.DB.copy()

def get_full_states_tree(depth):
    G = GraphVisualization()
    options = ['0']
    new_options = []

    for i in range(1,depth-1):
        for option in options:
            for j in range(4):
                new_option = option + str(j)
                new_options.append(new_option)
                G.addEdge(option, new_option)
                # print(option +","+ new_option)

        options = new_options
        new_options = []
    for option in options:
            new_option = option + str(0)
            new_options.append(new_option)
            G.addEdge(option, new_option)
            # print(option +","+ new_option)


    return G


# Defining a Class
class GraphVisualization:

    def __init__(self,arr=[]):

        # visual is a list which stores all
        # the set of edges that constitutes a
        # graph
        self.visual = []
        for x in arr:
            self.visual.append(x)
        self.G = nx.Graph()


    # addEdge function inputs the vertices of an
    # edge and appends it to the visual list
    def addEdge(self, a, b):
        temp = [a, b]
        self.visual.append(temp)
        self.G.add_edges_from([temp])


    # In visualize function G is an object of
    # class Graph given by networkx G.add_edges_from(visual)
    # creates a graph with a given list
    # nx.draw_networkx(G) - plots the graph
    # plt.show() - displays the graph
    def visualize_kamda(self):
        nx.draw_networkx(self.G,pos=nx.kamada_kawai_layout(self.G),node_size=10,with_labels=False)
        plt.show()

    def visualize(self):
        nx.draw_networkx(self.G,node_size=10,with_labels=False)
        plt.show()

# G = get_full_states_tree(5)
# G.visualize()