forked from PolicyEngine/policyengine-us-variable-leaves
-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
136 lines (101 loc) · 3.2 KB
/
main.py
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
122
123
124
125
126
127
128
129
130
131
132
133
134
135
from pyvis.network import Network
import subprocess
import argparse
import yaml
import webbrowser
import os
'''
run using "python main.py -v [variable_name]"
optional -s flag to specify a state
then open the nx.html file in the browser
'''
YEAR = 2024
def create_yaml_file(variable_name, state_code='CO'):
data = [
{
"name": "Scenario",
"period": YEAR,
"input": {
"state_code": state_code.upper()
},
"output": {variable_name: 0},
}
]
with open("test.yaml", "w") as file:
yaml.dump(data, file)
def run_policy_engine():
command = "policyengine-core test test.yaml -c policyengine_us -v"
result = subprocess.run(command, shell=True, capture_output=True, text=True)
return result.stdout
def nodes_and_edges(output: str):
lines = output.split("\n")
nodes = set()
leaf_nodes = set()
non_leaf_nodes = set()
edges = set()
parent_stack = []
for i, line in enumerate(lines):
if "=" not in line or f"<{YEAR}" not in line:
continue
variable = line.split("<")[0].strip()
current_indent = get_indent(line)
# Determine if the next line is more indented
next_line = lines[i + 1]
next_indent = get_indent(next_line)
nodes.add(variable)
if len(parent_stack):
# because of month variables, there can be variables nested under themselves
if parent_stack[-1] != variable:
edges.add((parent_stack[-1], variable))
if current_indent > next_indent:
parent_stack = parent_stack[:-(current_indent - next_indent)]
elif current_indent < next_indent:
parent_stack.append(variable)
if current_indent >= next_indent:
leaf_nodes.add(variable)
else:
non_leaf_nodes.add(variable)
leaf_nodes.difference_update(non_leaf_nodes)
return nodes, edges, leaf_nodes
def get_indent(line: str):
return int((len(line) - len(line.lstrip())) / 2)
parser = argparse.ArgumentParser()
parser.add_argument("-v", "--variable", help="variable name")
parser.add_argument("-s", "--state-code", help="set the state code", default='CO')
args = parser.parse_args()
create_yaml_file(args.variable, args.state_code)
output = run_policy_engine()
# Extract and display unique leaf nodes
nodes, edges, leaf_nodes = nodes_and_edges(output)
net = Network(
height="100vh", directed=True, select_menu=True, neighborhood_highlight=True
)
Network.set_options(
net,
'''
const options = {
"physics": {
"repulsion": {
"theta": 1,
"centralGravity": 0,
"springLength": 255,
"springConstant": 0.06,
"damping": 1,
"avoidOverlap": 1
},
"minVelocity": 0.75,
"solver": "repulsion"
}
}
''',
)
for node in nodes:
if node in leaf_nodes:
net.add_node(node, color="red")
else:
net.add_node(node)
for edge in edges:
net.add_edge(edge[0], edge[1])
HTML_FILE = 'nx.html'
net.write_html(HTML_FILE)
webbrowser.open('file://' + os.path.realpath(HTML_FILE))