generated from pitmonticone/LeanProject
-
Notifications
You must be signed in to change notification settings - Fork 67
/
Copy pathgenerate_image.py
executable file
·146 lines (121 loc) · 4.11 KB
/
generate_image.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
136
137
138
139
140
141
142
143
144
145
146
#!/usr/bin/env python3
"""
Example usage:
```sh
$ pip install pillow
$ python scripts/generate_image.py equational_theories/Subgraph.lean --close --filter
$ open equational_theories/Subgraph.png
$ python scripts/generate_image.py equational_theories/Generated/TrivialBruteforce/theorems/Apply.lean --close --filter --equations equational_theories/Equations/Basic.lean
$ python scripts/generate_image.py equational_theories/Generated/ equational_theories/Subgraph.lean --output equational_theories/Subgraph.png
```
"""
from process_implications import *
from PIL import Image # pip install pillow
import glob
############################################
def name_to_ind(name):
return int(name.removeprefix("Equation"))
UNKNOWN_COLOR = (0, 0, 0)
KNOWN_IMPLIES_COLOR = (0, 255, 0)
KNOWN_NOT_IMPLIES_COLOR = (255, 0, 0)
def append_dict(d, a, b):
if a in d:
d[a].append(b)
d[a].sort()
else:
d[a] = [b]
def close(known_implies):
# copilot wrote this
closure2 = {}
for a, b in known_implies:
append_dict(closure2, a, b)
new_pairs2 = {}
for a, b in known_implies:
append_dict(new_pairs2, b, a)
new_pairs = closure = set(known_implies)
while new_pairs:
print(f"current closure is size {len(closure)}")
new_pairs = {
(a, d) for a, b in new_pairs for c, d in known_implies if b == c
} - closure
print(f"found {len(new_pairs)} pairs for the closure")
closure |= new_pairs
return closure
def print_usage():
print(
"Usage: python process_implications.py <file_name.lean> [--close] [--filter] [--output <filename.png>] [--equations <equations.lean> ...]"
)
if __name__ == "__main__":
close_implies = False
filter_universe = False
equations_files = []
output_file = None
files = []
try:
i = 1
while i < len(argv):
current_arg = argv[i]
if "--close" == current_arg:
close_implies = True
i += 1
elif "--filter" == current_arg:
filter_universe = True
i += 1
elif "--equations" == current_arg:
equations_files.append(argv[i + 1])
i += 2
elif "--output" == current_arg:
output_file = argv[i + 1]
i += 2
else:
if os.path.isdir(current_arg):
for file_name in glob.glob(
current_arg + "/**/*.lean", recursive=True
):
files.append(file_name)
elif os.path.isfile(current_arg):
files.append(current_arg)
else:
assert False
i += 1
except Exception as e:
print(e)
print_usage()
exit(1)
if len(files) == 0:
print_usage()
exit(1)
if output_file is None:
output_file = files[0].removesuffix(".lean") + ".png"
print("Reading implications and contrimplications")
universe, known_implies, known_not_implies = parse_proofs_files(
equations_files, files
)
print(
f"Found {len(known_implies)} implications, and {len(known_not_implies)} contrimplications."
)
if close_implies:
print("Calculating closure")
known_implies = close(known_implies)
inds = {e: name_to_ind(e) for e in universe}
if filter_universe:
print("Filtering to universe")
inds = {
kv[0]: i for i, kv in enumerate(sorted(inds.items(), key=lambda kv: kv[1]))
}
min_ind = min(inds.values())
max_ind = max(inds.values())
size = max_ind - min_ind + 1
def ind(name):
return inds[name] - min_ind
print("Building picture")
data = [UNKNOWN_COLOR] * (size * size)
for known in known_implies:
a, b = known
data[ind(a) * size + ind(b)] = KNOWN_IMPLIES_COLOR
for known in known_not_implies:
a, b = known
data[ind(a) * size + ind(b)] = KNOWN_NOT_IMPLIES_COLOR
img = Image.new("RGB", (size, size))
img.putdata(data)
img.save(output_file)