X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=grids.py;h=5dad6f376529abf8514fb28d26693eb95bc65245;hb=f2ab5fd489adebe9b34ac825d39e41f13f287cb2;hp=ba09225212ee70c7b21ba3ea7a3b9f84dfb8421c;hpb=870d6808ac616b81cae00d9cb1f4de29bae23410;p=culture.git diff --git a/grids.py b/grids.py index ba09225..5dad6f3 100755 --- a/grids.py +++ b/grids.py @@ -37,10 +37,34 @@ class Grids(problem.Problem): max_nb_cached_chunks=None, chunk_size=None, nb_threads=-1, + tasks=None, ): self.colors = torch.tensor([c for _, c in self.named_colors]) self.height = 10 self.width = 10 + self.cache_rec_coo = {} + + all_tasks = [ + self.task_replace_color, + self.task_translate, + self.task_grow, + self.task_half_fill, + self.task_frame, + self.task_detect, + self.task_count, + self.task_trajectory, + self.task_bounce, + self.task_scale, + self.task_symbols, + self.task_isometry, + # self.task_path, + ] + + if tasks is None: + self.all_tasks = all_tasks + else: + self.all_tasks = [getattr(self, "task_" + t) for t in tasks.split(",")] + super().__init__(max_nb_cached_chunks, chunk_size, nb_threads) ###################################################################### @@ -199,41 +223,134 @@ class Grids(problem.Problem): def nb_token_values(self): return len(self.colors) - def rec_coo(self, nb_rec, min_height=3, min_width=3): - N = 10 + # @torch.compile + def rec_coo( + self, + nb_rec, + min_height=3, + min_width=3, + surface_max=None, + prevent_overlap=False, + ): + if surface_max is None: + surface_max = self.height * self.width // 2 + + signature = (nb_rec, min_height, min_width, surface_max) + + try: + return self.cache_rec_coo[signature].pop() + except IndexError: + pass + except KeyError: + pass + + N = 10000 while True: - i = torch.randint(self.height, (N, nb_rec, 2)).sort(dim=-1).values - j = torch.randint(self.width, (N, nb_rec, 2)).sort(dim=-1).values - if nb_rec == 2: - A_i1, A_i2, A_j1, A_j2 = i[:, 0, 0], i[:, 0, 1], j[:, 0, 0], j[:, 0, 1] - B_i1, B_i2, B_j1, B_j2 = i[:, 1, 0], i[:, 1, 1], j[:, 1, 0], j[:, 1, 1] - no_overlap = torch.logical_not( - (A_i1 > B_i2) & (A_i2 < B_i1) & (A_j1 > B_j1) & (A_j2 < B_j1) + while True: + i = torch.randint(self.height, (N * nb_rec, 2)).sort(dim=-1).values + j = torch.randint(self.width, (N * nb_rec, 2)).sort(dim=-1).values + + big_enough = ( + (i[:, 1] >= i[:, 0] + min_height) + & (j[:, 1] >= j[:, 0] + min_height) + & ((i[:, 1] - i[:, 0]) * (j[:, 1] - j[:, 0]) <= surface_max) ) - i, j = i[no_overlap], j[no_overlap] - elif nb_rec == 3: - A_i1, A_i2, A_j1, A_j2 = i[:, 0, 0], i[:, 0, 1], j[:, 0, 0], j[:, 0, 1] - B_i1, B_i2, B_j1, B_j2 = i[:, 1, 0], i[:, 1, 1], j[:, 1, 0], j[:, 1, 1] - C_i1, C_i2, C_j1, C_j2 = i[:, 2, 0], i[:, 2, 1], j[:, 2, 0], j[:, 2, 1] - no_overlap = ( - torch.logical_not( - (A_i1 > B_i2) & (A_i2 < B_i1) & (A_j1 > B_j1) & (A_j2 < B_j1) + + i, j = i[big_enough], j[big_enough] + + n = i.size(0) - i.size(0) % nb_rec + + if n > 0: + break + + i = i[:n].reshape(n // nb_rec, nb_rec, -1) + j = j[:n].reshape(n // nb_rec, nb_rec, -1) + + if prevent_overlap: + can_fit = ((i[:, :, 1] - i[:, :, 0]) * (j[:, :, 1] - j[:, :, 0])).sum( + dim=-1 + ) <= self.height * self.width + i, j = i[can_fit], j[can_fit] + if nb_rec == 2: + A_i1, A_i2, A_j1, A_j2 = ( + i[:, 0, 0], + i[:, 0, 1], + j[:, 0, 0], + j[:, 0, 1], ) - & torch.logical_not( - (A_i1 > C_i2) & (A_i2 < C_i1) & (A_j1 > C_j1) & (A_j2 < C_j1) + B_i1, B_i2, B_j1, B_j2 = ( + i[:, 1, 0], + i[:, 1, 1], + j[:, 1, 0], + j[:, 1, 1], ) - & torch.logical_not( - (B_i1 > C_i2) & (B_i2 < C_i1) & (B_j1 > C_j1) & (B_j2 < C_j1) + no_overlap = torch.logical_not( + (A_i1 >= B_i2) + & (A_i2 <= B_i1) + & (A_j1 >= B_j1) + & (A_j2 <= B_j1) ) - ) - i, j = (i[no_overlap], j[no_overlap]) - else: - assert nb_rec == 1 + i, j = i[no_overlap], j[no_overlap] + elif nb_rec == 3: + A_i1, A_i2, A_j1, A_j2 = ( + i[:, 0, 0], + i[:, 0, 1], + j[:, 0, 0], + j[:, 0, 1], + ) + B_i1, B_i2, B_j1, B_j2 = ( + i[:, 1, 0], + i[:, 1, 1], + j[:, 1, 0], + j[:, 1, 1], + ) + C_i1, C_i2, C_j1, C_j2 = ( + i[:, 2, 0], + i[:, 2, 1], + j[:, 2, 0], + j[:, 2, 1], + ) + no_overlap = ( + ( + (A_i1 >= B_i2) + | (A_i2 <= B_i1) + | (A_j1 >= B_j2) + | (A_j2 <= B_j1) + ) + & ( + (A_i1 >= C_i2) + | (A_i2 <= C_i1) + | (A_j1 >= C_j2) + | (A_j2 <= C_j1) + ) + & ( + (B_i1 >= C_i2) + | (B_i2 <= C_i1) + | (B_j1 >= C_j2) + | (B_j2 <= C_j1) + ) + ) + i, j = (i[no_overlap], j[no_overlap]) + else: + assert nb_rec == 1 if i.size(0) > 1: break - return [(i[0, k, 0], j[0, k, 0], i[0, k, 1], j[0, k, 1]) for k in range(nb_rec)] + self.cache_rec_coo[signature] = [ + [ + ( + i[n, k, 0].item(), + j[n, k, 0].item(), + i[n, k, 1].item(), + j[n, k, 1].item(), + ) + for k in range(nb_rec) + ] + for n in range(i.size(0)) + ] + + return self.cache_rec_coo[signature].pop() ###################################################################### @@ -242,7 +359,7 @@ class Grids(problem.Problem): nb_rec = 3 c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1 for X, f_X in [(A, f_A), (B, f_B)]: - r = self.rec_coo(nb_rec) + r = self.rec_coo(nb_rec, prevent_overlap=True) for n in range(nb_rec): i1, j1, i2, j2 = r[n] X[i1:i2, j1:j2] = c[n] @@ -250,12 +367,16 @@ class Grids(problem.Problem): # @torch.compile def task_translate(self, A, f_A, B, f_B): - di, dj = torch.randint(3, (2,)) - 1 + while True: + di, dj = torch.randint(3, (2,)) - 1 + if di.abs() + dj.abs() > 0: + break + nb_rec = 3 c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1 for X, f_X in [(A, f_A), (B, f_B)]: while True: - r = self.rec_coo(nb_rec) + r = self.rec_coo(nb_rec, prevent_overlap=True) i1, j1, i2, j2 = r[nb_rec - 1] if ( i1 + di >= 0 @@ -281,7 +402,7 @@ class Grids(problem.Problem): direction = torch.randint(2, (1,)) for X, f_X in [(A, f_A), (B, f_B)]: while True: - r = self.rec_coo(nb_rec) + r = self.rec_coo(nb_rec, prevent_overlap=True) i1, j1, i2, j2 = r[nb_rec - 1] if i1 + 3 < i2 and j1 + 3 < j2: break @@ -300,13 +421,13 @@ class Grids(problem.Problem): f_X[i1:i2, j1:j2] = c[n] # @torch.compile - def task_color_grow(self, A, f_A, B, f_B): + def task_half_fill(self, A, f_A, B, f_B): di, dj = torch.randint(2, (2,)) * 2 - 1 nb_rec = 3 c = torch.randperm(len(self.colors) - 1)[: 2 * nb_rec] + 1 direction = torch.randint(4, (1,)) for X, f_X in [(A, f_A), (B, f_B)]: - r = self.rec_coo(nb_rec) + r = self.rec_coo(nb_rec, prevent_overlap=True) for n in range(nb_rec): i1, j1, i2, j2 = r[n] X[i1:i2, j1:j2] = c[2 * n] @@ -346,20 +467,24 @@ class Grids(problem.Problem): nb_rec = 3 c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1 for X, f_X in [(A, f_A), (B, f_B)]: - r = self.rec_coo(nb_rec) + r = self.rec_coo(nb_rec, prevent_overlap=True) for n in range(nb_rec): i1, j1, i2, j2 = r[n] X[i1:i2, j1:j2] = c[n] - f_X[i1:i2, j1:j2] = c[n] if n == nb_rec - 1: - f_X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = 0 + f_X[i1:i2, j1] = c[n] + f_X[i1:i2, j2 - 1] = c[n] + f_X[i1, j1:j2] = c[n] + f_X[i2 - 1, j1:j2] = c[n] + else: + f_X[i1:i2, j1:j2] = c[n] # @torch.compile def task_detect(self, A, f_A, B, f_B): nb_rec = 3 c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1 for X, f_X in [(A, f_A), (B, f_B)]: - r = self.rec_coo(nb_rec) + r = self.rec_coo(nb_rec, prevent_overlap=True) for n in range(nb_rec): i1, j1, i2, j2 = r[n] X[i1:i2, j1:j2] = c[n] @@ -613,7 +738,7 @@ class Grids(problem.Problem): f_X[i[0] : i[0] + delta, j[0] : j[0] + delta] = c[q] # @torch.compile - def task_ortho(self, A, f_A, B, f_B): + def task_isometry(self, A, f_A, B, f_B): nb_rec = 3 di, dj = torch.randint(3, (2,)) - 1 o = torch.tensor([[0.0, 1.0], [-1.0, 0.0]]) @@ -667,9 +792,97 @@ class Grids(problem.Problem): ): break + def compute_distance(self, walls, goal_i, goal_j, start_i, start_j): + max_length = walls.numel() + dist = torch.full_like(walls, max_length) + + dist[goal_i, goal_j] = 0 + pred_dist = torch.empty_like(dist) + + while True: + pred_dist.copy_(dist) + d = ( + torch.cat( + ( + dist[None, 1:-1, 0:-2], + dist[None, 2:, 1:-1], + dist[None, 1:-1, 2:], + dist[None, 0:-2, 1:-1], + ), + 0, + ).min(dim=0)[0] + + 1 + ) + + dist[1:-1, 1:-1].minimum_(d) # = torch.min(dist[1:-1, 1:-1], d) + dist = walls * max_length + (1 - walls) * dist + + if dist[start_i, start_j] < max_length or dist.equal(pred_dist): + return dist * (1 - walls) + # @torch.compile - def task_islands(self, A, f_A, B, f_B): - pass + def task_path(self, A, f_A, B, f_B): + c = torch.randperm(len(self.colors) - 1)[:3] + 1 + dist = torch.empty(self.height + 2, self.width + 2) + for X, f_X in [(A, f_A), (B, f_B)]: + nb_rec = torch.randint(3, (1,)) + 1 + while True: + r = self.rec_coo(nb_rec, prevent_overlap=True) + X[...] = 0 + f_X[...] = 0 + for n in range(nb_rec): + i1, j1, i2, j2 = r[n] + X[i1:i2, j1:j2] = c[0] + f_X[i1:i2, j1:j2] = c[0] + while True: + i0, j0 = torch.randint(self.height, (1,)), torch.randint( + self.width, (1,) + ) + if X[i0, j0] == 0: + break + while True: + i1, j1 = torch.randint(self.height, (1,)), torch.randint( + self.width, (1,) + ) + if X[i1, j1] == 0: + break + dist[...] = 1 + dist[1:-1, 1:-1] = (X != 0).long() + dist[...] = self.compute_distance(dist, i1 + 1, j1 + 1, i0 + 1, j0 + 1) + if dist[i0 + 1, j0 + 1] >= 1 and dist[i0 + 1, j0 + 1] < self.height * 4: + break + + dist[1:-1, 1:-1] += (X != 0).long() * self.height * self.width + dist[0, :] = self.height * self.width + dist[-1, :] = self.height * self.width + dist[:, 0] = self.height * self.width + dist[:, -1] = self.height * self.width + # dist += torch.rand(dist.size()) + + i, j = i0 + 1, j0 + 1 + while i != i1 + 1 or j != j1 + 1: + f_X[i - 1, j - 1] = c[2] + r, s, t, u = ( + dist[i - 1, j], + dist[i, j - 1], + dist[i + 1, j], + dist[i, j + 1], + ) + m = min(r, s, t, u) + if r == m: + i = i - 1 + elif t == m: + i = i + 1 + elif s == m: + j = j - 1 + else: + j = j + 1 + + X[i0, j0] = c[2] + # f_X[i0, j0] = c[1] + + X[i1, j1] = c[1] + f_X[i1, j1] = c[1] # for X, f_X in [(A, f_A), (B, f_B)]: # n = torch.arange(self.height * self.width).reshape(self.height, self.width) @@ -679,24 +892,75 @@ class Grids(problem.Problem): # i,j=q%self.height,q//self.height # if - ###################################################################### + # @torch.compile + def task_puzzle(self, A, f_A, B, f_B): + S = 4 + i0, j0 = (self.height - S) // 2, (self.width - S) // 2 + c = torch.randperm(len(self.colors) - 1)[:4] + 1 + for X, f_X in [(A, f_A), (B, f_B)]: + while True: + f_X[...] = 0 + h = list(torch.randperm(c.size(0))) + n = torch.zeros(c.max() + 1) + for _ in range(2): + k = torch.randperm(S * S) + for q in k: + i, j = q % S + i0, q // S + j0 + if f_X[i, j] == 0: + r, s, t, u = ( + f_X[i - 1, j], + f_X[i, j - 1], + f_X[i + 1, j], + f_X[i, j + 1], + ) + r, s, t, u = torch.tensor([r, s, t, u])[torch.randperm(4)] + if r > 0 and n[r] < 6: + n[r] += 1 + f_X[i, j] = r + elif s > 0 and n[s] < 6: + n[s] += 1 + f_X[i, j] = s + elif t > 0 and n[t] < 6: + n[t] += 1 + f_X[i, j] = t + elif u > 0 and n[u] < 6: + n[u] += 1 + f_X[i, j] = u + else: + if len(h) > 0: + d = c[h.pop()] + n[d] += 1 + f_X[i, j] = d + + if n.sum() == S * S: + break - def all_tasks(self): - return [ - self.task_replace_color, - self.task_translate, - self.task_grow, - self.task_color_grow, - self.task_frame, - self.task_detect, - self.task_count, - self.task_trajectory, - self.task_bounce, - self.task_scale, - self.task_symbols, - self.task_ortho, - # self.task_islands, - ] + k = 0 + for d in range(4): + while True: + ii, jj = torch.randint(self.height, (1,)), torch.randint( + self.width, (1,) + ) + e = 0 + for i in range(S): + for j in range(S): + if ( + ii + i >= self.height + or jj + j >= self.width + or ( + f_X[i + i0, j + j0] == c[d] + and X[ii + i, jj + j] > 0 + ) + ): + e = 1 + if e == 0: + break + for i in range(S): + for j in range(S): + if f_X[i + i0, j + j0] == c[d]: + X[ii + i, jj + j] = c[d] + + ###################################################################### def trivial_prompts_and_answers(self, prompts, answers): S = self.height * self.width @@ -706,7 +970,7 @@ class Grids(problem.Problem): def generate_prompts_and_answers_(self, nb, tasks=None, progress_bar=False): if tasks is None: - tasks = self.all_tasks() + tasks = self.all_tasks S = self.height * self.width prompts = torch.zeros(nb, 3 * S + 2, dtype=torch.int64) @@ -752,6 +1016,15 @@ class Grids(problem.Problem): nrow, ) + def save_some_examples(self, result_dir): + nb, nrow = 72, 4 + for t in self.all_tasks: + print(t.__name__) + prompts, answers = self.generate_prompts_and_answers_(nb, tasks=[t]) + self.save_quizzes( + result_dir, t.__name__, prompts[:nb], answers[:nb], nrow=nrow + ) + ###################################################################### @@ -773,19 +1046,24 @@ if __name__ == "__main__": # exit(0) # if True: - # nb = 72 - - # for t in grids.all_tasks(): - # for t in [grids.task_count]: - # print(t.__name__) - # prompts, answers = grids.generate_prompts_and_answers_(nb, tasks=[t]) - # grids.save_quizzes("/tmp", t.__name__, prompts[:nb], answers[:nb], nrow=4) + nb, nrow = 72, 4 + # nb, nrow = 8, 2 + + # for t in grids.all_tasks: + for t in [ + grids.task_replace_color, + grids.task_frame, + ]: + print(t.__name__) + prompts, answers = grids.generate_prompts_and_answers_(nb, tasks=[t]) + grids.save_quizzes("/tmp", t.__name__, prompts[:nb], answers[:nb], nrow=nrow) - # exit(0) + exit(0) nb = 1000 - for t in grids.all_tasks(): + for t in grids.all_tasks: + # for t in [ grids.task_replace_color ]: #grids.all_tasks: start_time = time.perf_counter() prompts, answers = grids.generate_prompts_and_answers_(nb, tasks=[t]) delay = time.perf_counter() - start_time