X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;ds=sidebyside;f=grids.py;h=002a33ffe4d56bfe0459b19cf6897a97e55babd4;hb=2f87c91cf606a068de1450d198660de7e44cd356;hp=85d640dda9b3c6550cf7891037be1759a17a734e;hpb=428fd9169ecc3d03c9e8282d319682ddab0f098d;p=culture.git

diff --git a/grids.py b/grids.py
index 85d640d..002a33f 100755
--- a/grids.py
+++ b/grids.py
@@ -37,11 +37,34 @@ class Grids(problem.Problem):
         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_path,
+        ]
+
+        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)
 
     ######################################################################
@@ -344,7 +367,11 @@ class Grids(problem.Problem):
 
     # @torch.compile
     def task_translate(self, A, f_A, B, f_B):
-        di, dj = torch.randint(3, (2,)) - 1
+        while True:
+            di, dj = torch.randint(3, (2,)) - 1
+            if di.abs() + dj.abs() > 0:
+                break
+
         nb_rec = 3
         c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1
         for X, f_X in [(A, f_A), (B, f_B)]:
@@ -394,7 +421,7 @@ class Grids(problem.Problem):
                     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
@@ -711,7 +738,7 @@ class Grids(problem.Problem):
             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]])
@@ -765,9 +792,97 @@ class Grids(problem.Problem):
                 ):
                     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_islands(self, A, f_A, B, f_B):
-        pass
+    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]
 
     # for X, f_X in [(A, f_A), (B, f_B)]:
     # n = torch.arange(self.height * self.width).reshape(self.height, self.width)
@@ -777,24 +892,75 @@ class Grids(problem.Problem):
     # 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_islands,
-        ]
+            k = 0
+            for d in range(4):
+                while True:
+                    ii, jj = torch.randint(self.height, (1,)), torch.randint(
+                        self.width, (1,)
+                    )
+                    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 trivial_prompts_and_answers(self, prompts, answers):
         S = self.height * self.width
@@ -804,7 +970,7 @@ class Grids(problem.Problem):
 
     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)
@@ -830,7 +996,7 @@ class Grids(problem.Problem):
 
         return prompts.flatten(1), answers.flatten(1)
 
-    def save_quizzes(
+    def save_quiz_illustrations(
         self,
         result_dir,
         filename_prefix,
@@ -850,6 +1016,15 @@ class Grids(problem.Problem):
             nrow,
         )
 
+    def save_some_examples(self, result_dir):
+        nb, nrow = 72, 4
+        for t in self.all_tasks:
+            print(t.__name__)
+            prompts, answers = self.generate_prompts_and_answers_(nb, tasks=[t])
+            self.save_quiz_illustrations(
+                result_dir, t.__name__, prompts[:nb], answers[:nb], nrow=nrow
+            )
+
 
 ######################################################################
 
@@ -871,18 +1046,26 @@ if __name__ == "__main__":
     # exit(0)
 
     # if True:
-    nb = 72
-
-    for t in grids.all_tasks():
-        # for t in [grids.task_replace_color]:
+    nb, nrow = 72, 4
+    # nb, nrow = 8, 2
+
+    # for t in grids.all_tasks:
+    for t in [
+        grids.task_replace_color,
+        grids.task_frame,
+    ]:
         print(t.__name__)
         prompts, answers = grids.generate_prompts_and_answers_(nb, tasks=[t])
-        grids.save_quizzes("/tmp", t.__name__, prompts[:nb], answers[:nb], nrow=4)
+        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_replace_color ]: #grids.all_tasks:
         start_time = time.perf_counter()
         prompts, answers = grids.generate_prompts_and_answers_(nb, tasks=[t])
         delay = time.perf_counter() - start_time
@@ -894,7 +1077,7 @@ if __name__ == "__main__":
     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],