+ # @torch.compile
+ 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,))):
+ m = m @ o
+ if torch.rand(1) < 0.5:
+ m[0, :] = -m[0, :]
+
+ ci, cj = (self.height - 1) / 2, (self.width - 1) / 2
+
+ for X, f_X in [(A, f_A), (B, f_B)]:
+ while True:
+ X[...] = 0
+ f_X[...] = 0
+
+ c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1
+
+ for r in range(nb_rec):
+ while True:
+ i1, i2 = torch.randint(self.height - 2, (2,)) + 1
+ j1, j2 = torch.randint(self.width - 2, (2,)) + 1
+ if (
+ i2 >= i1
+ and j2 >= j1
+ and max(i2 - i1, j2 - j1) >= 2
+ and min(i2 - i1, j2 - j1) <= 3
+ ):
+ break
+ X[i1 : i2 + 1, j1 : j2 + 1] = c[r]
+
+ i1, j1, i2, j2 = i1 - ci, j1 - cj, i2 - ci, j2 - cj
+
+ i1, j1 = m[0, 0] * i1 + m[0, 1] * j1, m[1, 0] * i1 + m[1, 1] * j1
+ i2, j2 = m[0, 0] * i2 + m[0, 1] * j2, m[1, 0] * i2 + m[1, 1] * j2
+
+ i1, j1, i2, j2 = i1 + ci, j1 + cj, i2 + ci, j2 + cj
+ i1, i2 = i1.long() + di, i2.long() + di
+ j1, j2 = j1.long() + dj, j2.long() + dj
+ if i1 > i2:
+ i1, i2 = i2, i1
+ if j1 > j2:
+ j1, j2 = j2, j1
+
+ f_X[i1 : i2 + 1, j1 : j2 + 1] = c[r]
+
+ n = F.one_hot(X.flatten()).sum(dim=0)[1:]
+ if (
+ n.sum() > self.height * self.width // 4
+ and (n > 0).long().sum() == nb_rec
+ ):
+ 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_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]