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]),
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)
######################################################################
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)
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):
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):
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
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]
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]
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
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]
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, :]
):
break
- def compute_distance(self, walls, goal_i, goal_j, start_i, start_j):
+ def compute_distance(self, walls, goal_i, goal_j):
max_length = walls.numel()
dist = torch.full_like(walls, max_length)
while True:
pred_dist.copy_(dist)
- d = (
+ dist[1:-1, 1:-1] = (
torch.cat(
(
+ dist[None, 1:-1, 1:-1],
dist[None, 1:-1, 0:-2],
dist[None, 2:, 1:-1],
dist[None, 1:-1, 2:],
+ 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):
+ if dist.equal(pred_dist):
return dist * (1 - walls)
# @torch.compile
- def task_path(self, A, f_A, B, f_B):
+ def task_distance(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)
+ dist0 = torch.empty(self.height + 2, self.width + 2)
+ dist1 = 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
+ nb_rec = torch.randint(3, (1,)).item() + 1
while True:
r = self.rec_coo(nb_rec, prevent_overlap=True)
X[...] = 0
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,)
+ 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,)), torch.randint(
- self.width, (1,)
+ 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:
+ dist1[...] = 1
+ dist1[1:-1, 1:-1] = (X != 0).long()
+ dist1[...] = self.compute_distance(dist1, i1 + 1, j1 + 1)
+ if (
+ dist1[i0 + 1, j0 + 1] >= 1
+ and dist1[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
+ dist0[...] = 1
+ dist0[1:-1, 1:-1] = (X != 0).long()
+ dist0[...] = self.compute_distance(dist0, i0 + 1, j0 + 1)
- X[i0, j0] = c[2]
- # f_X[i0, j0] = c[1]
+ dist0 = dist0[1:-1, 1:-1]
+ dist1 = dist1[1:-1, 1:-1]
+
+ D = dist1[i0, j0]
+ for d in range(1, D):
+ M = (dist0 == d) & (dist1 == D - d)
+ f_X[...] = (1 - M) * f_X + M * c[1]
- X[i1, j1] = c[1]
- f_X[i1, j1] = c[1]
+ X[i0, j0] = c[2]
+ f_X[i0, j0] = c[2]
+ X[i1, j1] = c[2]
+ f_X[i1, j1] = c[2]
# for X, f_X in [(A, f_A), (B, f_B)]:
# n = torch.arange(self.height * self.width).reshape(self.height, self.width)
# 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_path,
- ]
+ 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
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)
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,
def save_some_examples(self, result_dir):
nb, nrow = 72, 4
- for t in self.all_tasks():
+ for t in self.all_tasks:
print(t.__name__)
prompts, answers = self.generate_prompts_and_answers_(nb, tasks=[t])
- self.save_quizzes(
+ self.save_quiz_illustrations(
result_dir, t.__name__, prompts[:nb], answers[:nb], nrow=nrow
)
nb, nrow = 72, 4
# nb, nrow = 8, 2
- # for t in grids.all_tasks():
- for t in [grids.task_path]:
+ # for t in grids.all_tasks:
+ for t in [grids.task_distance]:
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)
+ 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_distance]:
start_time = time.perf_counter()
prompts, answers = grids.generate_prompts_and_answers_(nb, tasks=[t])
delay = time.perf_counter() - start_time
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],