Update.
[picoclvr.git] / expr.py
1 #!/usr/bin/env python
2
3 import math, re
4
5 import torch, torchvision
6
7 from torch import nn
8 from torch.nn import functional as F
9
10
11 def random_var(nb_variables=None, variables=None):
12     if variables is None:
13         return chr(ord("A") + torch.randint(nb_variables, (1,)).item())
14     else:
15         l = list(variables)
16         return l[torch.randint(len(l), (1,)).item()]
17
18
19 def random_expr(variables, budget):
20     if budget <= 5:
21         op = torch.randint(2, (1,)).item()
22         if op == 0 and len(variables) > 0:
23             return random_var(variables=variables)
24         else:
25             return str(torch.randint(10, (1,)).item())
26     else:
27         op = torch.randint(4, (1,)).item()
28         if op == 0:
29             e = random_expr(variables, budget - 2)
30             if ("+" in e or "-" in e or "*" in e) and (e[0] != "(" or e[-1] != ")"):
31                 return "(" + e + ")"
32             else:
33                 return e
34         else:
35             b = 2 + torch.randint(budget - 5, (1,)).item()
36             e1 = random_expr(variables, b)
37             e2 = random_expr(variables, budget - b - 1)
38             if op == 1:
39                 return e1 + "+" + e2
40             elif op == 2:
41                 return e1 + "+" + e2
42             elif op == 3:
43                 return e1 + "*" + e2
44
45
46 def generate_program(nb_variables, length):
47     s = ""
48     variables = set()
49     while len(s) < length:
50         v = random_var(nb_variables=nb_variables)
51         s += v + "=" + random_expr(variables, budget=20) + ";"
52         variables.add(v)
53     return s, variables
54
55
56 def extract_results(seq):
57     f = lambda a: (a[0], -1 if a[1] == "" else int(a[1]))
58     results = [
59         dict([f(tuple(x.split(":"))) for x in re.findall("[A-Z]:[0-9]*", s)])
60         for s in seq
61     ]
62     return results
63
64
65 def generate_sequences(nb, nb_variables=5, length=20, randomize_length=False):
66     sequences = []
67     for n in range(nb):
68         result = None
69         while result == None or max(result.values()) > 100:
70             l = length
71             if l > 5 and randomize_length:
72                 l = 5 + torch.randint(l - 5, (1,)).item()
73             p, v = generate_program(nb_variables, l)
74             v = ", ".join(['"' + v + '": ' + v for v in v])
75             ldict = {}
76             exec(p + "result={" + v + "}", globals(), ldict)
77             result = ldict["result"]
78
79         k = list(result.keys())
80         k.sort()
81         sequences.append(p + " " + "".join([v + ":" + str(result[v]) + ";" for v in k]))
82
83     return sequences
84
85
86 if __name__ == "__main__":
87     import time
88
89     start_time = time.perf_counter()
90     sequences = generate_sequences(1000)
91     end_time = time.perf_counter()
92     for s in sequences[:10]:
93         print(s)
94     print(f"{len(sequences) / (end_time - start_time):.02f} samples per second")
95
96     print(extract_results(sequences[:10]))