X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;ds=inline;f=grids.py;h=e6519401d5cf3d134c2c9ce355b62845e3520712;hb=5b7022591f48382ec84b1dda17297b1ed15166d7;hp=85d640dda9b3c6550cf7891037be1759a17a734e;hpb=428fd9169ecc3d03c9e8282d319682ddab0f098d;p=culture.git diff --git a/grids.py b/grids.py index 85d640d..e651940 100755 --- a/grids.py +++ b/grids.py @@ -17,6 +17,92 @@ from torch.nn import functional as F import problem +def grow_islands(nb, height, width, nb_seeds, nb_iterations): + w = torch.empty(5, 1, 3, 3) + + w[0, 0] = torch.tensor( + [ + [1.0, 1.0, 1.0], + [1.0, 0.0, 1.0], + [1.0, 1.0, 1.0], + ] + ) + + w[1, 0] = torch.tensor( + [ + [-1.0, 1.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 0.0, 0.0], + ] + ) + + w[2, 0] = torch.tensor( + [ + [0.0, 1.0, -1.0], + [0.0, 0.0, 1.0], + [0.0, 0.0, 0.0], + ] + ) + + w[3, 0] = torch.tensor( + [ + [0.0, 0.0, 0.0], + [0.0, 0.0, 1.0], + [0.0, 1.0, -1.0], + ] + ) + + w[4, 0] = torch.tensor( + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [-1.0, 1.0, 0.0], + ] + ) + + Z = torch.zeros(nb, height, width) + U = Z.flatten(1) + + for _ in range(nb_seeds): + M = F.conv2d(Z[:, None, :, :], w, padding=1) + M = torch.cat([M[:, :1], M[:, 1:].min(dim=1, keepdim=True).values], dim=1) + M = ((M[:, 0] == 0) & (Z == 0)).long() + Q = (M.flatten(1).max(dim=1).values > 0).long()[:, None] + M = M * torch.rand(M.size()) + M = M.flatten(1) + M = F.one_hot(M.argmax(dim=1), num_classes=M.size(1)) + U += M * Q + + for _ in range(nb_iterations): + M = F.conv2d(Z[:, None, :, :], w, padding=1) + M = torch.cat([M[:, :1], M[:, 1:].min(dim=1, keepdim=True).values], dim=1) + M = ((M[:, 1] >= 0) & (Z == 0)).long() + Q = (M.flatten(1).max(dim=1).values > 0).long()[:, None] + M = M * torch.rand(M.size()) + M = M.flatten(1) + M = F.one_hot(M.argmax(dim=1), num_classes=M.size(1)) + U = Z.flatten(1) + U += M * Q + + M = Z.clone() + Z = Z * (torch.arange(Z.size(1) * Z.size(2)) + 1).reshape(1, Z.size(1), Z.size(2)) + + while True: + W = Z.clone() + Z = F.max_pool2d(Z, 3, 1, 1) * M + if Z.equal(W): + break + + Z = Z.long() + U = Z.flatten(1) + V = F.one_hot(U).max(dim=1).values + W = V.cumsum(dim=1) - V + N = torch.arange(Z.size(0))[:, None, None].expand_as(Z) + Z = W[N, Z] + + return Z + + class Grids(problem.Problem): named_colors = [ ("white", [255, 255, 255]), @@ -37,11 +123,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_islands, + ] + + 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) ###################################################################### @@ -344,7 +453,11 @@ 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)]: @@ -372,7 +485,7 @@ class Grids(problem.Problem): di, dj = torch.randint(2, (2,)) * 2 - 1 nb_rec = 3 c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1 - direction = torch.randint(2, (1,)) + direction = torch.randint(2, (1,)).item() for X, f_X in [(A, f_A), (B, f_B)]: while True: r = self.rec_coo(nb_rec, prevent_overlap=True) @@ -394,11 +507,11 @@ 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,)) + direction = torch.randint(4, (1,)).item() for X, f_X in [(A, f_A), (B, f_B)]: r = self.rec_coo(nb_rec, prevent_overlap=True) for n in range(nb_rec): @@ -501,61 +614,50 @@ class Grids(problem.Problem): return no, nq, nq_diag def task_count(self, A, f_A, B, f_B): - N = (torch.randint(4, (1,)) + 2).item() - c = torch.randperm(len(self.colors) - 1)[:N] + 1 - - for X, f_X in [(A, f_A), (B, f_B)]: - l_q = torch.randperm(self.height * self.width)[ - : self.height * self.width // 20 - ] - l_d = torch.randint(N, l_q.size()) - nb = torch.zeros(N, dtype=torch.int64) - - for q, e in zip(l_q, l_d): - d = c[e] - i, j = q % self.height, q // self.height - if ( - nb[e] < self.width - and X[max(0, i - 1) : i + 2, max(0, j - 1) : j + 2] == 0 - ).all(): - X[i, j] = d - nb[e] += 1 - - l_q = torch.randperm((self.height - 2) * (self.width - 2))[ - : self.height * self.width // 2 - ] - l_d = torch.randint(N, l_q.size()) - for q, e in zip(l_q, l_d): - d = c[e] - i, j = q % (self.height - 2) + 1, q // (self.height - 2) + 1 - a1, a2, a3 = X[i - 1, j - 1 : j + 2] - a8, a4 = X[i, j - 1], X[i, j + 1] - a7, a6, a5 = X[i + 1, j - 1 : j + 2] - if ( - X[i, j] == 0 - and nb[e] < self.width - and (a2 == 0 or a2 == d) - and (a4 == 0 or a4 == d) - and (a6 == 0 or a6 == d) - and (a8 == 0 or a8 == d) - and (a1 == 0 or a2 == d or a8 == d) - and (a3 == 0 or a4 == d or a2 == d) - and (a5 == 0 or a6 == d or a4 == d) - and (a7 == 0 or a8 == d or a6 == d) - ): - o = ( - (a2 != 0).long() - + (a4 != 0).long() - + (a6 != 0).long() - + (a8 != 0).long() + while True: + error = False + + N = torch.randint(5, (1,)).item() + 1 + c = torch.zeros(N + 1) + c[1:] = torch.randperm(len(self.colors) - 1)[:N] + 1 + + for X, f_X in [(A, f_A), (B, f_B)]: + if not hasattr(self, "cache_count") or len(self.cache_count) == 0: + self.cache_count = list( + grow_islands( + 1000, + self.height, + self.width, + nb_seeds=self.height * self.width // 8, + nb_iterations=self.height * self.width // 10, + ) ) - if o <= 1: - X[i, j] = d - nb[e] += 1 - o - for e in range(N): - for j in range(nb[e]): - f_X[e, j] = c[e] + X[...] = self.cache_count.pop() + + k = (X.max() + 1 + (c.size(0) - 1)).item() + V = torch.arange(k) // (c.size(0) - 1) + V = (V + torch.rand(V.size())).sort().indices[: X.max() + 1] % ( + c.size(0) - 1 + ) + 1 + V[0] = 0 + X[...] = c[V[X]] + + if F.one_hot(X.flatten()).max(dim=0).values.sum().item() == N + 1: + f_X[...] = 0 + for e in range(1, N + 1): + for j in range((X == c[e]).sum() + 1): + if j < self.width: + f_X[e - 1, j] = c[e] + else: + error = True + break + else: + error = True + break + + if not error: + break # @torch.compile def task_trajectory(self, A, f_A, B, f_B): @@ -563,7 +665,10 @@ class Grids(problem.Problem): for X, f_X in [(A, f_A), (B, f_B)]: while True: di, dj = torch.randint(7, (2,)) - 3 - i, j = torch.randint(self.height, (1,)), torch.randint(self.width, (1,)) + i, j = ( + torch.randint(self.height, (1,)).item(), + torch.randint(self.width, (1,)).item(), + ) if ( abs(di) + abs(dj) > 0 and i + 2 * di >= 0 @@ -604,8 +709,9 @@ class Grids(problem.Problem): X[...] = 0 for _ in range((self.height * self.width) // 10): - i, j = torch.randint(self.height, (1,)), torch.randint( - self.width, (1,) + i, j = ( + torch.randint(self.height, (1,)).item(), + torch.randint(self.width, (1,)).item(), ) X[i, j] = c[0] f_X[i, j] = c[0] @@ -615,7 +721,10 @@ class Grids(problem.Problem): if abs(di) + abs(dj) == 1: break - i, j = torch.randint(self.height, (1,)), torch.randint(self.width, (1,)) + i, j = ( + torch.randint(self.height, (1,)).item(), + torch.randint(self.width, (1,)).item(), + ) X[i, j] = c[1] f_X[i, j] = c[1] @@ -653,18 +762,21 @@ class Grids(problem.Problem): def task_scale(self, A, f_A, B, f_B): c = torch.randperm(len(self.colors) - 1)[:2] + 1 - i, j = torch.randint(self.height // 2, (1,)), torch.randint( - self.width // 2, (1,) + i, j = ( + torch.randint(self.height // 2, (1,)).item(), + torch.randint(self.width // 2, (1,)).item(), ) for X, f_X in [(A, f_A), (B, f_B)]: for _ in range(3): while True: - i1, j1 = torch.randint(self.height // 2 + 1, (1,)), torch.randint( - self.width // 2 + 1, (1,) + i1, j1 = ( + torch.randint(self.height // 2 + 1, (1,)).item(), + torch.randint(self.width // 2 + 1, (1,)).item(), ) - i2, j2 = torch.randint(self.height // 2 + 1, (1,)), torch.randint( - self.width // 2 + 1, (1,) + i2, j2 = ( + torch.randint(self.height // 2 + 1, (1,)).item(), + torch.randint(self.width // 2 + 1, (1,)).item(), ) if i1 < i2 and j1 < j2 and min(i2 - i1, j2 - j1) <= 3: break @@ -694,7 +806,7 @@ class Grids(problem.Problem): ai, aj = i.float().mean(), j.float().mean() - q = torch.randint(3, (1,)) + 1 + q = torch.randint(3, (1,)).item() + 1 X[i[0] + delta // 2 - 1, j[0] + delta // 2 - 1] = c[0] X[i[0] + delta // 2 - 1, j[0] + delta // 2 + 1] = c[0] @@ -711,12 +823,12 @@ 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]]) m = torch.eye(2) - for _ in range(torch.randint(4, (1,))): + for _ in range(torch.randint(4, (1,)).item()): m = m @ o if torch.rand(1) < 0.5: m[0, :] = -m[0, :] @@ -765,9 +877,99 @@ 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,)).item() + 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,)).item(), + torch.randint(self.width, (1,)).item(), + ) + if X[i0, j0] == 0: + break + while True: + i1, j1 = ( + torch.randint(self.height, (1,)).item(), + torch.randint(self.width, (1,)).item(), + ) + 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) @@ -777,24 +979,104 @@ 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,)).item(), + torch.randint(self.width, (1,)).item(), + ) + 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 task_islands(self, A, f_A, B, f_B): + c = torch.randperm(len(self.colors) - 1)[:2] + 1 + for X, f_X in [(A, f_A), (B, f_B)]: + if not hasattr(self, "cache_islands") or len(self.cache_islands) == 0: + self.cache_islands = list( + grow_islands( + 1000, + self.height, + self.width, + nb_seeds=self.height * self.width // 20, + nb_iterations=self.height * self.width // 2, + ) + ) + + A = self.cache_islands.pop() + + while True: + i, j = ( + torch.randint(self.height // 2, (1,)).item(), + torch.randint(self.width // 2, (1,)).item(), + ) + if A[i, j] > 0: + break + + X[...] = (A > 0) * c[0] + X[i, j] = c[1] + f_X[...] = (A == A[i, j]) * c[1] + ((A > 0) & (A != A[i, j])) * c[0] + + ###################################################################### def trivial_prompts_and_answers(self, prompts, answers): S = self.height * self.width @@ -804,7 +1086,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) @@ -825,12 +1107,12 @@ class Grids(problem.Problem): f_A = prompt[1 * (S + 1) : 1 * (S + 1) + S].view(self.height, self.width) B = prompt[2 * (S + 1) : 2 * (S + 1) + S].view(self.height, self.width) f_B = answer.view(self.height, self.width) - task = tasks[torch.randint(len(tasks), (1,))] + task = tasks[torch.randint(len(tasks), (1,)).item()] task(A, f_A, B, f_B) return prompts.flatten(1), answers.flatten(1) - def save_quizzes( + def save_quiz_illustrations( self, result_dir, filename_prefix, @@ -850,6 +1132,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_quiz_illustrations( + result_dir, t.__name__, prompts[:nb], answers[:nb], nrow=nrow + ) + ###################################################################### @@ -871,18 +1162,23 @@ if __name__ == "__main__": # exit(0) # if True: - nb = 72 + nb, nrow = 72, 4 + # nb, nrow = 8, 2 - for t in grids.all_tasks(): - # for t in [grids.task_replace_color]: + # for t in grids.all_tasks: + for t in [grids.task_islands]: 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) + grids.save_quiz_illustrations( + "/tmp", t.__name__, prompts[:nb], answers[:nb], nrow=nrow + ) + + # exit(0) nb = 1000 - for t in grids.all_tasks(): - # for t in [ grids.task_replace_color ]: #grids.all_tasks(): + # for t in grids.all_tasks: + for t in [grids.task_islands]: start_time = time.perf_counter() prompts, answers = grids.generate_prompts_and_answers_(nb, tasks=[t]) delay = time.perf_counter() - start_time @@ -894,7 +1190,7 @@ if __name__ == "__main__": predicted_prompts = m * (torch.randint(2, (prompts.size(0),)) * 2 - 1) predicted_answers = (1 - m) * (torch.randint(2, (prompts.size(0),)) * 2 - 1) - grids.save_quizzes( + grids.save_quiz_illustrations( "/tmp", "test", prompts[:nb],