X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=grids.py;h=20a964b8baeaff1cef3c2b02e89d9ee92a06a8a7;hb=a86dff174205c38d8e90d0d89ea399a6afb36359;hp=ba09225212ee70c7b21ba3ea7a3b9f84dfb8421c;hpb=870d6808ac616b81cae00d9cb1f4de29bae23410;p=culture.git

diff --git a/grids.py b/grids.py
index ba09225..20a964b 100755
--- a/grids.py
+++ b/grids.py
@@ -41,6 +41,7 @@ class Grids(problem.Problem):
         self.colors = torch.tensor([c for _, c in self.named_colors])
         self.height = 10
         self.width = 10
+        self.cache_rec_coo = {}
         super().__init__(max_nb_cached_chunks, chunk_size, nb_threads)
 
     ######################################################################
@@ -199,41 +200,134 @@ class Grids(problem.Problem):
     def nb_token_values(self):
         return len(self.colors)
 
-    def rec_coo(self, nb_rec, min_height=3, min_width=3):
-        N = 10
+    # @torch.compile
+    def rec_coo(
+        self,
+        nb_rec,
+        min_height=3,
+        min_width=3,
+        surface_max=None,
+        prevent_overlap=False,
+    ):
+        if surface_max is None:
+            surface_max = self.height * self.width // 2
+
+        signature = (nb_rec, min_height, min_width, surface_max)
+
+        try:
+            return self.cache_rec_coo[signature].pop()
+        except IndexError:
+            pass
+        except KeyError:
+            pass
+
+        N = 10000
         while True:
-            i = torch.randint(self.height, (N, nb_rec, 2)).sort(dim=-1).values
-            j = torch.randint(self.width, (N, nb_rec, 2)).sort(dim=-1).values
-            if nb_rec == 2:
-                A_i1, A_i2, A_j1, A_j2 = i[:, 0, 0], i[:, 0, 1], j[:, 0, 0], j[:, 0, 1]
-                B_i1, B_i2, B_j1, B_j2 = i[:, 1, 0], i[:, 1, 1], j[:, 1, 0], j[:, 1, 1]
-                no_overlap = torch.logical_not(
-                    (A_i1 > B_i2) & (A_i2 < B_i1) & (A_j1 > B_j1) & (A_j2 < B_j1)
+            while True:
+                i = torch.randint(self.height, (N * nb_rec, 2)).sort(dim=-1).values
+                j = torch.randint(self.width, (N * nb_rec, 2)).sort(dim=-1).values
+
+                big_enough = (
+                    (i[:, 1] >= i[:, 0] + min_height)
+                    & (j[:, 1] >= j[:, 0] + min_height)
+                    & ((i[:, 1] - i[:, 0]) * (j[:, 1] - j[:, 0]) <= surface_max)
                 )
-                i, j = i[no_overlap], j[no_overlap]
-            elif nb_rec == 3:
-                A_i1, A_i2, A_j1, A_j2 = i[:, 0, 0], i[:, 0, 1], j[:, 0, 0], j[:, 0, 1]
-                B_i1, B_i2, B_j1, B_j2 = i[:, 1, 0], i[:, 1, 1], j[:, 1, 0], j[:, 1, 1]
-                C_i1, C_i2, C_j1, C_j2 = i[:, 2, 0], i[:, 2, 1], j[:, 2, 0], j[:, 2, 1]
-                no_overlap = (
-                    torch.logical_not(
-                        (A_i1 > B_i2) & (A_i2 < B_i1) & (A_j1 > B_j1) & (A_j2 < B_j1)
+
+                i, j = i[big_enough], j[big_enough]
+
+                n = i.size(0) - i.size(0) % nb_rec
+
+                if n > 0:
+                    break
+
+            i = i[:n].reshape(n // nb_rec, nb_rec, -1)
+            j = j[:n].reshape(n // nb_rec, nb_rec, -1)
+
+            if prevent_overlap:
+                can_fit = ((i[:, :, 1] - i[:, :, 0]) * (j[:, :, 1] - j[:, :, 0])).sum(
+                    dim=-1
+                ) <= self.height * self.width
+                i, j = i[can_fit], j[can_fit]
+                if nb_rec == 2:
+                    A_i1, A_i2, A_j1, A_j2 = (
+                        i[:, 0, 0],
+                        i[:, 0, 1],
+                        j[:, 0, 0],
+                        j[:, 0, 1],
                     )
-                    & torch.logical_not(
-                        (A_i1 > C_i2) & (A_i2 < C_i1) & (A_j1 > C_j1) & (A_j2 < C_j1)
+                    B_i1, B_i2, B_j1, B_j2 = (
+                        i[:, 1, 0],
+                        i[:, 1, 1],
+                        j[:, 1, 0],
+                        j[:, 1, 1],
                     )
-                    & torch.logical_not(
-                        (B_i1 > C_i2) & (B_i2 < C_i1) & (B_j1 > C_j1) & (B_j2 < C_j1)
+                    no_overlap = torch.logical_not(
+                        (A_i1 >= B_i2)
+                        & (A_i2 <= B_i1)
+                        & (A_j1 >= B_j1)
+                        & (A_j2 <= B_j1)
                     )
-                )
-                i, j = (i[no_overlap], j[no_overlap])
-            else:
-                assert nb_rec == 1
+                    i, j = i[no_overlap], j[no_overlap]
+                elif nb_rec == 3:
+                    A_i1, A_i2, A_j1, A_j2 = (
+                        i[:, 0, 0],
+                        i[:, 0, 1],
+                        j[:, 0, 0],
+                        j[:, 0, 1],
+                    )
+                    B_i1, B_i2, B_j1, B_j2 = (
+                        i[:, 1, 0],
+                        i[:, 1, 1],
+                        j[:, 1, 0],
+                        j[:, 1, 1],
+                    )
+                    C_i1, C_i2, C_j1, C_j2 = (
+                        i[:, 2, 0],
+                        i[:, 2, 1],
+                        j[:, 2, 0],
+                        j[:, 2, 1],
+                    )
+                    no_overlap = (
+                        (
+                            (A_i1 >= B_i2)
+                            | (A_i2 <= B_i1)
+                            | (A_j1 >= B_j2)
+                            | (A_j2 <= B_j1)
+                        )
+                        & (
+                            (A_i1 >= C_i2)
+                            | (A_i2 <= C_i1)
+                            | (A_j1 >= C_j2)
+                            | (A_j2 <= C_j1)
+                        )
+                        & (
+                            (B_i1 >= C_i2)
+                            | (B_i2 <= C_i1)
+                            | (B_j1 >= C_j2)
+                            | (B_j2 <= C_j1)
+                        )
+                    )
+                    i, j = (i[no_overlap], j[no_overlap])
+                else:
+                    assert nb_rec == 1
 
             if i.size(0) > 1:
                 break
 
-        return [(i[0, k, 0], j[0, k, 0], i[0, k, 1], j[0, k, 1]) for k in range(nb_rec)]
+        self.cache_rec_coo[signature] = [
+            [
+                (
+                    i[n, k, 0].item(),
+                    j[n, k, 0].item(),
+                    i[n, k, 1].item(),
+                    j[n, k, 1].item(),
+                )
+                for k in range(nb_rec)
+            ]
+            for n in range(i.size(0))
+        ]
+
+        return self.cache_rec_coo[signature].pop()
 
     ######################################################################
 
@@ -242,7 +336,7 @@ class Grids(problem.Problem):
         nb_rec = 3
         c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
         for X, f_X in [(A, f_A), (B, f_B)]:
-            r = self.rec_coo(nb_rec)
+            r = self.rec_coo(nb_rec, prevent_overlap=True)
             for n in range(nb_rec):
                 i1, j1, i2, j2 = r[n]
                 X[i1:i2, j1:j2] = c[n]
@@ -250,12 +344,16 @@ 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)]:
             while True:
-                r = self.rec_coo(nb_rec)
+                r = self.rec_coo(nb_rec, prevent_overlap=True)
                 i1, j1, i2, j2 = r[nb_rec - 1]
                 if (
                     i1 + di >= 0
@@ -281,7 +379,7 @@ class Grids(problem.Problem):
         direction = torch.randint(2, (1,))
         for X, f_X in [(A, f_A), (B, f_B)]:
             while True:
-                r = self.rec_coo(nb_rec)
+                r = self.rec_coo(nb_rec, prevent_overlap=True)
                 i1, j1, i2, j2 = r[nb_rec - 1]
                 if i1 + 3 < i2 and j1 + 3 < j2:
                     break
@@ -306,7 +404,7 @@ class Grids(problem.Problem):
         c = torch.randperm(len(self.colors) - 1)[: 2 * nb_rec] + 1
         direction = torch.randint(4, (1,))
         for X, f_X in [(A, f_A), (B, f_B)]:
-            r = self.rec_coo(nb_rec)
+            r = self.rec_coo(nb_rec, prevent_overlap=True)
             for n in range(nb_rec):
                 i1, j1, i2, j2 = r[n]
                 X[i1:i2, j1:j2] = c[2 * n]
@@ -346,20 +444,24 @@ class Grids(problem.Problem):
         nb_rec = 3
         c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
         for X, f_X in [(A, f_A), (B, f_B)]:
-            r = self.rec_coo(nb_rec)
+            r = self.rec_coo(nb_rec, prevent_overlap=True)
             for n in range(nb_rec):
                 i1, j1, i2, j2 = r[n]
                 X[i1:i2, j1:j2] = c[n]
-                f_X[i1:i2, j1:j2] = c[n]
                 if n == nb_rec - 1:
-                    f_X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = 0
+                    f_X[i1:i2, j1] = c[n]
+                    f_X[i1:i2, j2 - 1] = c[n]
+                    f_X[i1, j1:j2] = c[n]
+                    f_X[i2 - 1, j1:j2] = c[n]
+                else:
+                    f_X[i1:i2, j1:j2] = c[n]
 
     # @torch.compile
     def task_detect(self, A, f_A, B, f_B):
         nb_rec = 3
         c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
         for X, f_X in [(A, f_A), (B, f_B)]:
-            r = self.rec_coo(nb_rec)
+            r = self.rec_coo(nb_rec, prevent_overlap=True)
             for n in range(nb_rec):
                 i1, j1, i2, j2 = r[n]
                 X[i1:i2, j1:j2] = c[n]
@@ -667,9 +769,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)
@@ -695,7 +885,7 @@ class Grids(problem.Problem):
             self.task_scale,
             self.task_symbols,
             self.task_ortho,
-            # self.task_islands,
+            #            self.task_path,
         ]
 
     def trivial_prompts_and_answers(self, prompts, answers):
@@ -752,6 +942,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_quizzes(
+                result_dir, t.__name__, prompts[:nb], answers[:nb], nrow=nrow
+            )
+
 
 ######################################################################
 
@@ -773,19 +972,21 @@ if __name__ == "__main__":
     # exit(0)
 
     # if True:
-    # nb = 72
+    nb, nrow = 72, 4
+    # nb, nrow = 8, 2
 
     # for t in grids.all_tasks():
-    # for t in [grids.task_count]:
-    # 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)
+    for t in [grids.task_path]:
+        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)
 
     # exit(0)
 
     nb = 1000
 
     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