grids.py

   1 #!/usr/bin/env python
   2
   3 # Any copyright is dedicated to the Public Domain.
   4 # https://creativecommons.org/publicdomain/zero/1.0/
   5
   6 # Written by Francois Fleuret <francois@fleuret.org>
   7
   8 import math, sys, tqdm, os, warnings
   9
  10 import torch, torchvision
  11
  12 from torch import nn
  13 from torch.nn import functional as F
  14
  15 ######################################################################
  16
  17 import problem
  18
  19
  20 class Grids(problem.Problem):
  21     named_colors = [
  22         ("white", [255, 255, 255]),
  23         ("red", [255, 0, 0]),
  24         ("green", [0, 192, 0]),
  25         ("blue", [0, 0, 255]),
  26         ("yellow", [255, 224, 0]),
  27         ("cyan", [0, 255, 255]),
  28         ("violet", [224, 128, 255]),
  29         ("lightgreen", [192, 255, 192]),
  30         ("brown", [165, 42, 42]),
  31         ("lightblue", [192, 192, 255]),
  32         ("gray", [128, 128, 128]),
  33     ]
  34
  35     def __init__(
  36         self,
  37         max_nb_cached_chunks=None,
  38         chunk_size=None,
  39         nb_threads=-1,
  40     ):
  41         self.colors = torch.tensor([c for _, c in self.named_colors])
  42         self.height = 10
  43         self.width = 10
  44         super().__init__(max_nb_cached_chunks, chunk_size, nb_threads)
  45
  46     ######################################################################
  47
  48     def frame2img(self, x, scale=15):
  49         x = x.reshape(x.size(0), self.height, -1)
  50         m = torch.logical_and(x >= 0, x < self.nb_token_values()).long()
  51         x = self.colors[x * m].permute(0, 3, 1, 2)
  52         s = x.shape
  53         x = x[:, :, :, None, :, None].expand(-1, -1, -1, scale, -1, scale)
  54         x = x.reshape(s[0], s[1], s[2] * scale, s[3] * scale)
  55
  56         x[:, :, :, torch.arange(0, x.size(3), scale)] = 0
  57         x[:, :, torch.arange(0, x.size(2), scale), :] = 0
  58         x = x[:, :, 1:, 1:]
  59
  60         for n in range(m.size(0)):
  61             for i in range(m.size(1)):
  62                 for j in range(m.size(2)):
  63                     if m[n, i, j] == 0:
  64                         for k in range(2, scale - 2):
  65                             for l in [0, 1]:
  66                                 x[n, :, i * scale + k, j * scale + k - l] = 0
  67                                 x[
  68                                     n, :, i * scale + scale - 1 - k, j * scale + k - l
  69                                 ] = 0
  70
  71         return x
  72
  73     def save_image(
  74         self,
  75         result_dir,
  76         filename,
  77         prompts,
  78         answers,
  79         predicted_prompts=None,
  80         predicted_answers=None,
  81         nrow=4,
  82         margin=8,
  83     ):
  84         S = self.height * self.width
  85         As = prompts[:, 0 * (S + 1) : 0 * (S + 1) + S].view(-1, self.height, self.width)
  86         f_As = prompts[:, 1 * (S + 1) : 1 * (S + 1) + S].view(
  87             -1, self.height, self.width
  88         )
  89         Bs = prompts[:, 2 * (S + 1) : 2 * (S + 1) + S].view(-1, self.height, self.width)
  90         prompts = torch.cat([As, f_As, Bs], dim=2)
  91         answers = answers.reshape(answers.size(0), self.height, self.width)
  92
  93         if predicted_prompts is None:
  94             predicted_prompts = 255
  95
  96         if predicted_answers is None:
  97             predicted_answers = 255
  98
  99         def add_frame(x, c, margin, bottom=False):
 100             if bottom:
 101                 h, w, di, dj = x.size(2) + margin, x.size(3), 0, 0
 102             else:
 103                 h, w, di, dj = (
 104                     x.size(2) + 2 * margin,
 105                     x.size(3) + 2 * margin,
 106                     margin,
 107                     margin,
 108                 )
 109
 110             y = x.new_full((x.size(0), x.size(1), h, w), 0)
 111
 112             if type(c) is int:
 113                 y[...] = c
 114             else:
 115                 c = c.long()[:, None]
 116                 c = (
 117                     (1 - ((c == 1).long() + (c == 0).long() + (c == -1).long()))
 118                     * torch.tensor([64, 64, 64])
 119                     + (c == 1).long() * torch.tensor([0, 255, 0])
 120                     + (c == 0).long() * torch.tensor([255, 255, 255])
 121                     + (c == -1).long() * torch.tensor([255, 0, 0])
 122                 )
 123                 y[...] = c[:, :, None, None]
 124
 125             y[:, :, di : di + x.size(2), dj : dj + x.size(3)] = x
 126
 127             return y
 128
 129         img_prompts = torch.cat(
 130             [
 131                 add_frame(
 132                     add_frame(self.frame2img(x), c=0, margin=1),
 133                     c=predicted_prompts,
 134                     margin=margin,
 135                 )
 136                 for x in prompts.to("cpu").split(split_size=self.width, dim=2)
 137             ],
 138             dim=3,
 139         )
 140
 141         h = img_prompts.size(2)
 142         img_answers = add_frame(
 143             add_frame(self.frame2img(answers.to("cpu")), c=0, margin=1),
 144             c=predicted_answers,
 145             margin=margin,
 146         )
 147
 148         separator_size = 2 * margin
 149
 150         separator = img_prompts.new_full(
 151             (
 152                 img_prompts.size(0),
 153                 img_prompts.size(1),
 154                 img_prompts.size(2),
 155                 separator_size,
 156             ),
 157             255,
 158         )
 159
 160         marker = img_prompts.new_full(
 161             (
 162                 img_prompts.size(0),
 163                 img_prompts.size(1),
 164                 img_prompts.size(2),
 165                 separator_size,
 166             ),
 167             255,
 168         )
 169
 170         # marker[:, :, 0] = 0
 171         # marker[:, :, h - 1] = 0
 172
 173         for k in range(1, 2 * separator_size - 8):
 174             i = k - (separator_size - 4)
 175             j = separator_size - 5 - abs(i)
 176             marker[:, :, h // 2 - 1 + i, 2 + j] = 0
 177             marker[:, :, h // 2 - 1 + i + 1, 2 + j] = 0
 178
 179         img = torch.cat(
 180             [
 181                 img_prompts,
 182                 marker,
 183                 img_answers,
 184             ],
 185             dim=3,
 186         )
 187
 188         image_name = os.path.join(result_dir, filename)
 189         torchvision.utils.save_image(
 190             img.float() / 255.0,
 191             image_name,
 192             nrow=nrow,
 193             padding=margin * 4,
 194             pad_value=1.0,
 195         )
 196
 197     ######################################################################
 198
 199     def nb_token_values(self):
 200         return len(self.colors)
 201
 202     def rec_coo(self, nb_rec, min_height=3, min_width=3):
 203         N = 10
 204         while True:
 205             i = torch.randint(self.height, (N, nb_rec, 2)).sort(dim=-1).values
 206             j = torch.randint(self.width, (N, nb_rec, 2)).sort(dim=-1).values
 207             if nb_rec == 2:
 208                 A_i1, A_i2, A_j1, A_j2 = i[:, 0, 0], i[:, 0, 1], j[:, 0, 0], j[:, 0, 1]
 209                 B_i1, B_i2, B_j1, B_j2 = i[:, 1, 0], i[:, 1, 1], j[:, 1, 0], j[:, 1, 1]
 210                 no_overlap = torch.logical_not(
 211                     (A_i1 > B_i2) & (A_i2 < B_i1) & (A_j1 > B_j1) & (A_j2 < B_j1)
 212                 )
 213                 i, j = i[no_overlap], j[no_overlap]
 214             elif nb_rec == 3:
 215                 A_i1, A_i2, A_j1, A_j2 = i[:, 0, 0], i[:, 0, 1], j[:, 0, 0], j[:, 0, 1]
 216                 B_i1, B_i2, B_j1, B_j2 = i[:, 1, 0], i[:, 1, 1], j[:, 1, 0], j[:, 1, 1]
 217                 C_i1, C_i2, C_j1, C_j2 = i[:, 2, 0], i[:, 2, 1], j[:, 2, 0], j[:, 2, 1]
 218                 no_overlap = (
 219                     torch.logical_not(
 220                         (A_i1 > B_i2) & (A_i2 < B_i1) & (A_j1 > B_j1) & (A_j2 < B_j1)
 221                     )
 222                     & torch.logical_not(
 223                         (A_i1 > C_i2) & (A_i2 < C_i1) & (A_j1 > C_j1) & (A_j2 < C_j1)
 224                     )
 225                     & torch.logical_not(
 226                         (B_i1 > C_i2) & (B_i2 < C_i1) & (B_j1 > C_j1) & (B_j2 < C_j1)
 227                     )
 228                 )
 229                 i, j = (i[no_overlap], j[no_overlap])
 230             else:
 231                 assert nb_rec == 1
 232
 233             if i.size(0) > 1:
 234                 break
 235
 236         return [(i[0, k, 0], j[0, k, 0], i[0, k, 1], j[0, k, 1]) for k in range(nb_rec)]
 237
 238     ######################################################################
 239
 240     # @torch.compile
 241     def task_replace_color(self, A, f_A, B, f_B):
 242         nb_rec = 3
 243         c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
 244         for X, f_X in [(A, f_A), (B, f_B)]:
 245             r = self.rec_coo(nb_rec)
 246             for n in range(nb_rec):
 247                 i1, j1, i2, j2 = r[n]
 248                 X[i1:i2, j1:j2] = c[n]
 249                 f_X[i1:i2, j1:j2] = c[n if n > 0 else -1]
 250
 251     # @torch.compile
 252     def task_translate(self, A, f_A, B, f_B):
 253         di, dj = torch.randint(3, (2,)) - 1
 254         nb_rec = 3
 255         c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1
 256         for X, f_X in [(A, f_A), (B, f_B)]:
 257             while True:
 258                 r = self.rec_coo(nb_rec)
 259                 i1, j1, i2, j2 = r[nb_rec - 1]
 260                 if (
 261                     i1 + di >= 0
 262                     and i2 + di < X.size(0)
 263                     and j1 + dj >= 0
 264                     and j2 + dj < X.size(1)
 265                 ):
 266                     break
 267
 268             for n in range(nb_rec):
 269                 i1, j1, i2, j2 = r[n]
 270                 X[i1:i2, j1:j2] = c[n]
 271                 if n == nb_rec - 1:
 272                     f_X[i1 + di : i2 + di, j1 + dj : j2 + dj] = c[n]
 273                 else:
 274                     f_X[i1:i2, j1:j2] = c[n]
 275
 276     # @torch.compile
 277     def task_grow(self, A, f_A, B, f_B):
 278         di, dj = torch.randint(2, (2,)) * 2 - 1
 279         nb_rec = 3
 280         c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1
 281         direction = torch.randint(2, (1,))
 282         for X, f_X in [(A, f_A), (B, f_B)]:
 283             while True:
 284                 r = self.rec_coo(nb_rec)
 285                 i1, j1, i2, j2 = r[nb_rec - 1]
 286                 if i1 + 3 < i2 and j1 + 3 < j2:
 287                     break
 288
 289             for n in range(nb_rec):
 290                 i1, j1, i2, j2 = r[n]
 291                 if n == nb_rec - 1:
 292                     if direction == 0:
 293                         X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = c[n]
 294                         f_X[i1:i2, j1:j2] = c[n]
 295                     else:
 296                         X[i1:i2, j1:j2] = c[n]
 297                         f_X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = c[n]
 298                 else:
 299                     X[i1:i2, j1:j2] = c[n]
 300                     f_X[i1:i2, j1:j2] = c[n]
 301
 302     # @torch.compile
 303     def task_color_grow(self, A, f_A, B, f_B):
 304         di, dj = torch.randint(2, (2,)) * 2 - 1
 305         nb_rec = 3
 306         c = torch.randperm(len(self.colors) - 1)[: 2 * nb_rec] + 1
 307         direction = torch.randint(4, (1,))
 308         for X, f_X in [(A, f_A), (B, f_B)]:
 309             r = self.rec_coo(nb_rec)
 310             for n in range(nb_rec):
 311                 i1, j1, i2, j2 = r[n]
 312                 X[i1:i2, j1:j2] = c[2 * n]
 313                 f_X[i1:i2, j1:j2] = c[2 * n]
 314                 # Not my proudest moment
 315                 if direction == 0:
 316                     i = (i1 + i2) // 2
 317                     X[i : i + 1, j1:j2] = c[2 * n + 1]
 318                     if n == nb_rec - 1:
 319                         f_X[i:i2, j1:j2] = c[2 * n + 1]
 320                     else:
 321                         f_X[i : i + 1, j1:j2] = c[2 * n + 1]
 322                 elif direction == 1:
 323                     i = (i1 + i2 - 1) // 2
 324                     X[i : i + 1, j1:j2] = c[2 * n + 1]
 325                     if n == nb_rec - 1:
 326                         f_X[i1 : i + 1, j1:j2] = c[2 * n + 1]
 327                     else:
 328                         f_X[i : i + 1, j1:j2] = c[2 * n + 1]
 329                 elif direction == 2:
 330                     j = (j1 + j2) // 2
 331                     X[i1:i2, j : j + 1] = c[2 * n + 1]
 332                     if n == nb_rec - 1:
 333                         f_X[i1:i2, j:j2] = c[2 * n + 1]
 334                     else:
 335                         f_X[i1:i2, j : j + 1] = c[2 * n + 1]
 336                 elif direction == 3:
 337                     j = (j1 + j2 - 1) // 2
 338                     X[i1:i2, j : j + 1] = c[2 * n + 1]
 339                     if n == nb_rec - 1:
 340                         f_X[i1:i2, j1 : j + 1] = c[2 * n + 1]
 341                     else:
 342                         f_X[i1:i2, j : j + 1] = c[2 * n + 1]
 343
 344     # @torch.compile
 345     def task_frame(self, A, f_A, B, f_B):
 346         nb_rec = 3
 347         c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
 348         for X, f_X in [(A, f_A), (B, f_B)]:
 349             r = self.rec_coo(nb_rec)
 350             for n in range(nb_rec):
 351                 i1, j1, i2, j2 = r[n]
 352                 X[i1:i2, j1:j2] = c[n]
 353                 f_X[i1:i2, j1:j2] = c[n]
 354                 if n == nb_rec - 1:
 355                     f_X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = 0
 356
 357     # @torch.compile
 358     def task_detect(self, A, f_A, B, f_B):
 359         nb_rec = 3
 360         c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
 361         for X, f_X in [(A, f_A), (B, f_B)]:
 362             r = self.rec_coo(nb_rec)
 363             for n in range(nb_rec):
 364                 i1, j1, i2, j2 = r[n]
 365                 X[i1:i2, j1:j2] = c[n]
 366                 if n < nb_rec - 1:
 367                     f_X[i1, j1] = c[-1]
 368
 369     # @torch.compile
 370     def contact(self, X, i, j, q):
 371         nq, nq_diag = 0, 0
 372         no = 0
 373
 374         for ii, jj in [
 375             (i - 1, j - 1),
 376             (i - 1, j),
 377             (i - 1, j + 1),
 378             (i, j - 1),
 379             (i, j + 1),
 380             (i + 1, j - 1),
 381             (i + 1, j),
 382             (i + 1, j + 1),
 383         ]:
 384             if ii >= 0 and ii < self.height and jj >= 0 and jj < self.width:
 385                 if X[ii, jj] != 0 and X[ii, jj] != q:
 386                     no += 1
 387
 388         for ii, jj in [
 389             (i - 1, j - 1),
 390             (i - 1, j + 1),
 391             (i + 1, j - 1),
 392             (i + 1, j + 1),
 393         ]:
 394             if ii >= 0 and ii < self.height and jj >= 0 and jj < self.width:
 395                 if X[ii, jj] == q and X[i, jj] != q and X[ii, j] != q:
 396                     nq_diag += 1
 397
 398         for ii, jj in [(i - 1, j), (i, j - 1), (i, j + 1), (i + 1, j)]:
 399             if ii >= 0 and ii < self.height and jj >= 0 and jj < self.width:
 400                 if X[ii, jj] == q:
 401                     nq += 1
 402
 403         return no, nq, nq_diag
 404
 405     def task_count(self, A, f_A, B, f_B):
 406         N = (torch.randint(4, (1,)) + 2).item()
 407         c = torch.randperm(len(self.colors) - 1)[:N] + 1
 408
 409         for X, f_X in [(A, f_A), (B, f_B)]:
 410             l_q = torch.randperm(self.height * self.width)[
 411                 : self.height * self.width // 20
 412             ]
 413             l_d = torch.randint(N, l_q.size())
 414             nb = torch.zeros(N, dtype=torch.int64)
 415
 416             for q, e in zip(l_q, l_d):
 417                 d = c[e]
 418                 i, j = q % self.height, q // self.height
 419                 if (
 420                     nb[e] < self.width
 421                     and X[max(0, i - 1) : i + 2, max(0, j - 1) : j + 2] == 0
 422                 ).all():
 423                     X[i, j] = d
 424                     nb[e] += 1
 425
 426             l_q = torch.randperm((self.height - 2) * (self.width - 2))[
 427                 : self.height * self.width // 2
 428             ]
 429             l_d = torch.randint(N, l_q.size())
 430             for q, e in zip(l_q, l_d):
 431                 d = c[e]
 432                 i, j = q % (self.height - 2) + 1, q // (self.height - 2) + 1
 433                 a1, a2, a3 = X[i - 1, j - 1 : j + 2]
 434                 a8, a4 = X[i, j - 1], X[i, j + 1]
 435                 a7, a6, a5 = X[i + 1, j - 1 : j + 2]
 436                 if (
 437                     X[i, j] == 0
 438                     and nb[e] < self.width
 439                     and (a2 == 0 or a2 == d)
 440                     and (a4 == 0 or a4 == d)
 441                     and (a6 == 0 or a6 == d)
 442                     and (a8 == 0 or a8 == d)
 443                     and (a1 == 0 or a2 == d or a8 == d)
 444                     and (a3 == 0 or a4 == d or a2 == d)
 445                     and (a5 == 0 or a6 == d or a4 == d)
 446                     and (a7 == 0 or a8 == d or a6 == d)
 447                 ):
 448                     o = (
 449                         (a2 != 0).long()
 450                         + (a4 != 0).long()
 451                         + (a6 != 0).long()
 452                         + (a8 != 0).long()
 453                     )
 454                     if o <= 1:
 455                         X[i, j] = d
 456                         nb[e] += 1 - o
 457
 458             for e in range(N):
 459                 for j in range(nb[e]):
 460                     f_X[e, j] = c[e]
 461
 462     # @torch.compile
 463     def task_trajectory(self, A, f_A, B, f_B):
 464         c = torch.randperm(len(self.colors) - 1)[:2] + 1
 465         for X, f_X in [(A, f_A), (B, f_B)]:
 466             while True:
 467                 di, dj = torch.randint(7, (2,)) - 3
 468                 i, j = torch.randint(self.height, (1,)), torch.randint(self.width, (1,))
 469                 if (
 470                     abs(di) + abs(dj) > 0
 471                     and i + 2 * di >= 0
 472                     and i + 2 * di < self.height
 473                     and j + 2 * dj >= 0
 474                     and j + 2 * dj < self.width
 475                 ):
 476                     break
 477
 478             k = 0
 479             while (
 480                 i + k * di >= 0
 481                 and i + k * di < self.height
 482                 and j + k * dj >= 0
 483                 and j + k * dj < self.width
 484             ):
 485                 if k < 2:
 486                     X[i + k * di, j + k * dj] = c[k]
 487                 f_X[i + k * di, j + k * dj] = c[min(k, 1)]
 488                 k += 1
 489
 490     # @torch.compile
 491     def task_bounce(self, A, f_A, B, f_B):
 492         c = torch.randperm(len(self.colors) - 1)[:3] + 1
 493         for X, f_X in [(A, f_A), (B, f_B)]:
 494             # @torch.compile
 495             def free(i, j):
 496                 return (
 497                     i >= 0
 498                     and i < self.height
 499                     and j >= 0
 500                     and j < self.width
 501                     and f_X[i, j] == 0
 502                 )
 503
 504             while True:
 505                 f_X[...] = 0
 506                 X[...] = 0
 507
 508                 for _ in range((self.height * self.width) // 10):
 509                     i, j = torch.randint(self.height, (1,)), torch.randint(
 510                         self.width, (1,)
 511                     )
 512                     X[i, j] = c[0]
 513                     f_X[i, j] = c[0]
 514
 515                 while True:
 516                     di, dj = torch.randint(7, (2,)) - 3
 517                     if abs(di) + abs(dj) == 1:
 518                         break
 519
 520                 i, j = torch.randint(self.height, (1,)), torch.randint(self.width, (1,))
 521
 522                 X[i, j] = c[1]
 523                 f_X[i, j] = c[1]
 524                 l = 0
 525
 526                 while True:
 527                     l += 1
 528                     if free(i + di, j + dj):
 529                         pass
 530                     elif free(i - dj, j + di):
 531                         di, dj = -dj, di
 532                         if free(i + dj, j - di):
 533                             if torch.rand(1) < 0.5:
 534                                 di, dj = -di, -dj
 535                     elif free(i + dj, j - di):
 536                         di, dj = dj, -di
 537                     else:
 538                         break
 539
 540                     i, j = i + di, j + dj
 541                     f_X[i, j] = c[2]
 542                     if l <= 1:
 543                         X[i, j] = c[2]
 544
 545                     if l >= self.width:
 546                         break
 547
 548                 f_X[i, j] = c[1]
 549                 X[i, j] = c[1]
 550
 551                 if l > 3:
 552                     break
 553
 554     # @torch.compile
 555     def task_scale(self, A, f_A, B, f_B):
 556         c = torch.randperm(len(self.colors) - 1)[:2] + 1
 557
 558         i, j = torch.randint(self.height // 2, (1,)), torch.randint(
 559             self.width // 2, (1,)
 560         )
 561
 562         for X, f_X in [(A, f_A), (B, f_B)]:
 563             for _ in range(3):
 564                 while True:
 565                     i1, j1 = torch.randint(self.height // 2 + 1, (1,)), torch.randint(
 566                         self.width // 2 + 1, (1,)
 567                     )
 568                     i2, j2 = torch.randint(self.height // 2 + 1, (1,)), torch.randint(
 569                         self.width // 2 + 1, (1,)
 570                     )
 571                     if i1 < i2 and j1 < j2 and min(i2 - i1, j2 - j1) <= 3:
 572                         break
 573                 X[i + i1 : i + i2, j + j1 : j + j2] = c[0]
 574                 f_X[2 * i1 : 2 * i2, 2 * j1 : 2 * j2] = c[0]
 575
 576             X[i, j] = c[1]
 577             f_X[0:2, 0:2] = c[1]
 578
 579     # @torch.compile
 580     def task_symbols(self, A, f_A, B, f_B):
 581         nb_rec = 4
 582         c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
 583         delta = 3
 584         for X, f_X in [(A, f_A), (B, f_B)]:
 585             while True:
 586                 i, j = torch.randint(self.height - delta + 1, (nb_rec,)), torch.randint(
 587                     self.width - delta + 1, (nb_rec,)
 588                 )
 589                 d = (i[None, :] - i[:, None]).abs().max((j[None, :] - j[:, None]).abs())
 590                 d.fill_diagonal_(delta + 1)
 591                 if d.min() > delta:
 592                     break
 593
 594             for k in range(1, nb_rec):
 595                 X[i[k] : i[k] + delta, j[k] : j[k] + delta] = c[k]
 596
 597             ai, aj = i.float().mean(), j.float().mean()
 598
 599             q = torch.randint(3, (1,)) + 1
 600
 601             X[i[0] + delta // 2 - 1, j[0] + delta // 2 - 1] = c[0]
 602             X[i[0] + delta // 2 - 1, j[0] + delta // 2 + 1] = c[0]
 603             X[i[0] + delta // 2 + 1, j[0] + delta // 2 - 1] = c[0]
 604             X[i[0] + delta // 2 + 1, j[0] + delta // 2 + 1] = c[0]
 605
 606             assert i[q] != ai and j[q] != aj
 607
 608             X[
 609                 i[0] + delta // 2 + (i[q] - ai).sign().long(),
 610                 j[0] + delta // 2 + (j[q] - aj).sign().long(),
 611             ] = c[nb_rec]
 612
 613             f_X[i[0] : i[0] + delta, j[0] : j[0] + delta] = c[q]
 614
 615     # @torch.compile
 616     def task_ortho(self, A, f_A, B, f_B):
 617         nb_rec = 3
 618         di, dj = torch.randint(3, (2,)) - 1
 619         o = torch.tensor([[0.0, 1.0], [-1.0, 0.0]])
 620         m = torch.eye(2)
 621         for _ in range(torch.randint(4, (1,))):
 622             m = m @ o
 623         if torch.rand(1) < 0.5:
 624             m[0, :] = -m[0, :]
 625
 626         ci, cj = (self.height - 1) / 2, (self.width - 1) / 2
 627
 628         for X, f_X in [(A, f_A), (B, f_B)]:
 629             while True:
 630                 X[...] = 0
 631                 f_X[...] = 0
 632
 633                 c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1
 634
 635                 for r in range(nb_rec):
 636                     while True:
 637                         i1, i2 = torch.randint(self.height - 2, (2,)) + 1
 638                         j1, j2 = torch.randint(self.width - 2, (2,)) + 1
 639                         if (
 640                             i2 >= i1
 641                             and j2 >= j1
 642                             and max(i2 - i1, j2 - j1) >= 2
 643                             and min(i2 - i1, j2 - j1) <= 3
 644                         ):
 645                             break
 646                     X[i1 : i2 + 1, j1 : j2 + 1] = c[r]
 647
 648                     i1, j1, i2, j2 = i1 - ci, j1 - cj, i2 - ci, j2 - cj
 649
 650                     i1, j1 = m[0, 0] * i1 + m[0, 1] * j1, m[1, 0] * i1 + m[1, 1] * j1
 651                     i2, j2 = m[0, 0] * i2 + m[0, 1] * j2, m[1, 0] * i2 + m[1, 1] * j2
 652
 653                     i1, j1, i2, j2 = i1 + ci, j1 + cj, i2 + ci, j2 + cj
 654                     i1, i2 = i1.long() + di, i2.long() + di
 655                     j1, j2 = j1.long() + dj, j2.long() + dj
 656                     if i1 > i2:
 657                         i1, i2 = i2, i1
 658                     if j1 > j2:
 659                         j1, j2 = j2, j1
 660
 661                     f_X[i1 : i2 + 1, j1 : j2 + 1] = c[r]
 662
 663                 n = F.one_hot(X.flatten()).sum(dim=0)[1:]
 664                 if (
 665                     n.sum() > self.height * self.width // 4
 666                     and (n > 0).long().sum() == nb_rec
 667                 ):
 668                     break
 669
 670     # @torch.compile
 671     def task_islands(self, A, f_A, B, f_B):
 672         pass
 673
 674     # for X, f_X in [(A, f_A), (B, f_B)]:
 675     # n = torch.arange(self.height * self.width).reshape(self.height, self.width)
 676     # k = torch.randperm(self.height * self.width)
 677     # X[...]=-1
 678     # for q in k:
 679     # i,j=q%self.height,q//self.height
 680     # if
 681
 682     ######################################################################
 683
 684     def all_tasks(self):
 685         return [
 686             self.task_replace_color,
 687             self.task_translate,
 688             self.task_grow,
 689             self.task_color_grow,
 690             self.task_frame,
 691             self.task_detect,
 692             self.task_count,
 693             self.task_trajectory,
 694             self.task_bounce,
 695             self.task_scale,
 696             self.task_symbols,
 697             self.task_ortho,
 698             # self.task_islands,
 699         ]
 700
 701     def trivial_prompts_and_answers(self, prompts, answers):
 702         S = self.height * self.width
 703         Bs = prompts[:, 2 * (S + 1) : 2 * (S + 1) + S]
 704         f_Bs = answers
 705         return (Bs == f_Bs).long().min(dim=-1).values > 0
 706
 707     def generate_prompts_and_answers_(self, nb, tasks=None, progress_bar=False):
 708         if tasks is None:
 709             tasks = self.all_tasks()
 710
 711         S = self.height * self.width
 712         prompts = torch.zeros(nb, 3 * S + 2, dtype=torch.int64)
 713         answers = torch.zeros(nb, S, dtype=torch.int64)
 714
 715         bunch = zip(prompts, answers)
 716
 717         if progress_bar:
 718             bunch = tqdm.tqdm(
 719                 bunch,
 720                 dynamic_ncols=True,
 721                 desc="world generation",
 722                 total=prompts.size(0),
 723             )
 724
 725         for prompt, answer in bunch:
 726             A = prompt[0 * (S + 1) : 0 * (S + 1) + S].view(self.height, self.width)
 727             f_A = prompt[1 * (S + 1) : 1 * (S + 1) + S].view(self.height, self.width)
 728             B = prompt[2 * (S + 1) : 2 * (S + 1) + S].view(self.height, self.width)
 729             f_B = answer.view(self.height, self.width)
 730             task = tasks[torch.randint(len(tasks), (1,))]
 731             task(A, f_A, B, f_B)
 732
 733         return prompts.flatten(1), answers.flatten(1)
 734
 735     def save_quizzes(
 736         self,
 737         result_dir,
 738         filename_prefix,
 739         prompts,
 740         answers,
 741         predicted_prompts=None,
 742         predicted_answers=None,
 743         nrow=4,
 744     ):
 745         self.save_image(
 746             result_dir,
 747             filename_prefix + ".png",
 748             prompts,
 749             answers,
 750             predicted_prompts,
 751             predicted_answers,
 752             nrow,
 753         )
 754
 755
 756 ######################################################################
 757
 758 if __name__ == "__main__":
 759     import time
 760
 761     # grids = Grids(max_nb_cached_chunks=5, chunk_size=100, nb_threads=4)
 762     grids = Grids()
 763
 764     # nb = 1000
 765     # grids = problem.MultiThreadProblem(
 766     # grids, max_nb_cached_chunks=50, chunk_size=100, nb_threads=1
 767     # )
 768     #    time.sleep(10)
 769     # start_time = time.perf_counter()
 770     # prompts, answers = grids.generate_prompts_and_answers(nb)
 771     # delay = time.perf_counter() - start_time
 772     # print(f"{prompts.size(0)/delay:02f} seq/s")
 773     # exit(0)
 774
 775     # if True:
 776     # nb = 72
 777
 778     # for t in grids.all_tasks():
 779     # for t in [grids.task_count]:
 780     # print(t.__name__)
 781     # prompts, answers = grids.generate_prompts_and_answers_(nb, tasks=[t])
 782     # grids.save_quizzes("/tmp", t.__name__, prompts[:nb], answers[:nb], nrow=4)
 783
 784     # exit(0)
 785
 786     nb = 1000
 787
 788     for t in grids.all_tasks():
 789         start_time = time.perf_counter()
 790         prompts, answers = grids.generate_prompts_and_answers_(nb, tasks=[t])
 791         delay = time.perf_counter() - start_time
 792         print(f"{t.__name__} {prompts.size(0)/delay:02f} seq/s")
 793
 794     exit(0)
 795
 796     m = torch.randint(2, (prompts.size(0),))
 797     predicted_prompts = m * (torch.randint(2, (prompts.size(0),)) * 2 - 1)
 798     predicted_answers = (1 - m) * (torch.randint(2, (prompts.size(0),)) * 2 - 1)
 799
 800     grids.save_quizzes(
 801         "/tmp",
 802         "test",
 803         prompts[:nb],
 804         answers[:nb],
 805         # You can add a bool to put a frame around the predicted parts
 806         predicted_prompts[:nb],
 807         predicted_answers[:nb],
 808     )