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