Update.

[culture.git] / grids.py

diff --git a/grids.py b/grids.py
index 8d144cf..eea8c6c 100755 (executable)
--- a/grids.py
+++ b/grids.py
@@ -67,26 +67,31 @@ def grow_islands(nb, height, width, nb_seeds, 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[:, 0] == 0) & (Z == 0)).long()
         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))
         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
+        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()
 
     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)
         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
+        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))
 
 
     M = Z.clone()
     Z = Z * (torch.arange(Z.size(1) * Z.size(2)) + 1).reshape(1, Z.size(1), Z.size(2))
 

-    for _ in range(100):
+    while True:
+        W = Z.clone()

         Z = F.max_pool2d(Z, 3, 1, 1) * M
         Z = F.max_pool2d(Z, 3, 1, 1) * M

+        if Z.equal(W):
+            break

 
     Z = Z.long()
     U = Z.flatten(1)
 
     Z = Z.long()
     U = Z.flatten(1)


@@ -138,7 +143,7 @@ class Grids(problem.Problem):
             self.task_scale,
             self.task_symbols,
             self.task_isometry,
             self.task_scale,
             self.task_symbols,
             self.task_isometry,

-            #            self.task_path,
+            #            self.task_islands,

         ]
 
         if tasks is None:
         ]
 
         if tasks is None:


@@ -609,61 +614,50 @@ class Grids(problem.Problem):
         return no, nq, nq_diag
 
     def task_count(self, A, f_A, B, f_B):
         return no, nq, nq_diag
 
     def task_count(self, A, f_A, B, f_B):

-        N = torch.randint(4, (1,)).item() + 2
-        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):
 
     # @torch.compile
     def task_trajectory(self, A, f_A, B, f_B):


@@ -883,7 +877,7 @@ class Grids(problem.Problem):
                 ):
                     break
 
                 ):
                     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)
 
         max_length = walls.numel()
         dist = torch.full_like(walls, max_length)
 


@@ -892,9 +886,10 @@ class Grids(problem.Problem):
 
         while True:
             pred_dist.copy_(dist)
 
         while True:
             pred_dist.copy_(dist)

-            d = (
+            dist[1:-1, 1:-1] = (

                 torch.cat(
                     (
                 torch.cat(
                     (

+                        dist[None, 1:-1, 1:-1],

                         dist[None, 1:-1, 0:-2],
                         dist[None, 2:, 1:-1],
                         dist[None, 1:-1, 2:],
                         dist[None, 1:-1, 0:-2],
                         dist[None, 2:, 1:-1],
                         dist[None, 1:-1, 2:],


@@ -905,16 +900,16 @@ class Grids(problem.Problem):
                 + 1
             )
 
                 + 1
             )
 

-            dist[1:-1, 1:-1].minimum_(d)  # = torch.min(dist[1:-1, 1:-1], d)

             dist = walls * max_length + (1 - walls) * dist
 
             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
                 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
         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,)).item() + 1
             while True:
         for X, f_X in [(A, f_A), (B, f_B)]:
             nb_rec = torch.randint(3, (1,)).item() + 1
             while True:


@@ -939,43 +934,31 @@ class Grids(problem.Problem):
                     )
                     if X[i1, j1] == 0:
                         break
                     )
                     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
 
                     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]

 
 

-            X[i1, j1] = c[1]
-            f_X[i1, j1] = c[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[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)
 
     # for X, f_X in [(A, f_A), (B, f_B)]:
     # n = torch.arange(self.height * self.width).reshape(self.height, self.width)


@@ -1057,65 +1040,30 @@ class Grids(problem.Problem):
     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)]:
     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)]:

-            while True:
-                k = torch.randperm(self.height * self.width)
-                Z = torch.zeros(self.height + 2, self.width + 2)
-
-                i0, j0 = (
-                    torch.randint(self.height, (1,)).item() + 1,
-                    torch.randint(self.width, (1,)).item() + 1,
+            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,
+                    )

                 )
 
                 )
 

-                Z[i0 - 1 : i0 + 2, j0 - 1 : j0 + 2] = 1
-
-                nb = 9
-
-                for q in k:
-                    i, j = q % self.height + 1, q // self.height + 1
-
-                    if Z[i, j] == 0:
-                        r, s, t, u, v, w, x, y = (
-                            Z[i - 1, j],
-                            Z[i - 1, j + 1],
-                            Z[i, j + 1],
-                            Z[i + 1, j + 1],
-                            Z[i + 1, j],
-                            Z[i + 1, j - 1],
-                            Z[i, j - 1],
-                            Z[i - 1, j - 1],
-                        )
-
-                        if (
-                            (nb < 16 or r + s + t + u + v + w + x + y > 0)
-                            and (s == 0 or r + t > 0)
-                            and (u == 0 or t + v > 0)
-                            and (w == 0 or x + v > 0)
-                            and (y == 0 or x + r > 0)
-                        ):
-                            # if r+s+t+u+v+w+x+y==0:
-                            Z[i, j] = 1
-                            nb += 1
-
-                    if nb == self.height * self.width // 2:
-                        break
-
-                if nb == self.height * self.width // 2:
-                    break
+            A = self.cache_islands.pop()

 
 

-            M = Z.clone()
-            Z[i0, j0] = 2
-            X[...] = (Z[1:-1, 1:-1] == 1) * c[0] + (Z[1:-1, 1:-1] == 2) * c[1]
-
-            for _ in range(self.height + self.width):
-                Z[1:-1, 1:-1] = Z[1:-1, 1:-1].maximum(
-                    torch.maximum(
-                        torch.maximum(Z[0:-2, 1:-1], Z[2:, 1:-1]),
-                        torch.maximum(Z[1:-1, 0:-2], Z[1:-1, 2:]),
-                    )
+            while True:
+                i, j = (
+                    torch.randint(self.height // 2, (1,)).item(),
+                    torch.randint(self.width // 2, (1,)).item(),

                 )
                 )

-                Z *= M
+                if A[i, j] > 0:
+                    break

 
 

-            f_X[...] = (Z[1:-1, 1:-1] == 1) * c[0] + (Z[1:-1, 1:-1] == 2) * c[1]
+            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]

 
     ######################################################################
 
 
     ######################################################################
 


@@ -1207,19 +1155,19 @@ if __name__ == "__main__":
     # nb, nrow = 8, 2
 
     # for t in grids.all_tasks:
     # nb, nrow = 8, 2
 
     # for t in grids.all_tasks:

-    for t in [grids.task_count]:
+    for t in [grids.task_distance]:

         print(t.__name__)
         prompts, answers = grids.generate_prompts_and_answers_(nb, tasks=[t])
         grids.save_quiz_illustrations(
             "/tmp", t.__name__, prompts[:nb], answers[:nb], nrow=nrow
         )
 
         print(t.__name__)
         prompts, answers = grids.generate_prompts_and_answers_(nb, tasks=[t])
         grids.save_quiz_illustrations(
             "/tmp", t.__name__, prompts[:nb], answers[:nb], nrow=nrow
         )
 

-    exit(0)
+    # exit(0)

 
     nb = 1000
 
     # for t in grids.all_tasks:
 
     nb = 1000
 
     # for t in grids.all_tasks:

-    for t in [grids.task_islands]:
+    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
         start_time = time.perf_counter()
         prompts, answers = grids.generate_prompts_and_answers_(nb, tasks=[t])
         delay = time.perf_counter() - start_time