close to a solution for 2

json_tests.py

@@ -10,7 +10,7 @@ from task import Task
 class TaskTester(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
-        cls.task_dict = load_task()#"data/training/0b148d64.json")
+        cls.task_dict = load_task("data/training/d4f3cd78.json")
         cls.task = Task(cls.task_dict)
         # Example of accessing input/output grids for the first example
         cls.output_grid = cls.task.train_examples[0].output_grid.grid
@@ -29,12 +29,13 @@ class TaskTester(unittest.TestCase):
         self.assertTrue((out_grid == self.output_grid).all())
 
     def test_task_class(self):
-        solution = "solution_1"
+        solution = "solution_2"
         res = self.task.try_solution(solution)
         self.assertTrue(res)
 
     def test_task_empty(self):
         res = self.task.try_solution("empty_solution")
         self.assertFalse(res)
+
 if __name__ == '__main__':
     unittest.main()

main.py

@@ -144,6 +144,11 @@ def colour_names2idx(colour_names):
     colour_to_idx = {colour: idx for idx, colour in zip(df.int, df.name)}
     # Vectorize the mapping function
     vectorized_mapping = np.vectorize(colour_to_idx.get)
+
+    incorrect_colours = np.unique(colour_names[~np.isin(colour_names, df.name)])
+    if len(incorrect_colours) > 0:
+        raise IndexError(f"Incorrect colour names: {'|'.join(incorrect_colours)}\n"
+                         f"Must be one of: {','.join(df.name)}")
     # Apply the mapping to the 2D array
     arr_idx = vectorized_mapping(colour_names)
     return arr_idx

prolog/solution_2.pl

+all_rows(Rs):-
+    setof(R, C^Colour^input_colour(R,C,Colour), Rs).
+
+all_cols(Cs):-
+    setof(C, R^Colour^input_colour(R,C,Colour), Cs).
+
+row(R):-
+    all_rows(Rs),
+    member(R,Rs).
+
+column(C):-
+    all_cols(Cs),
+    member(C,Cs).
+
+nrow(NR):-
+    all_rows(Rs),
+    length(Rs, NR).
+
+ncol(NC):-
+    all_cols(Cs),
+    length(Cs, NC).
+
+
+min(A, B, B):-
+    A < B.
+
+min(B, A, B):-
+    A >= B.
+min([Min], Min).
+
+min([H|T], Min):-
+    min(T, MinT),
+    min(H, MinT, Min).
+
+
+
+left_wall(C):-
+    input_colour(_,C,grey),
+    input_colour(_,C2,grey),
+    C2 > C.
+
+right_wall(C):-
+    input_colour(_,C,grey),
+    input_colour(_,C2,grey),
+    C2 < C.
+
+top_wall(R):-
+    input_colour(R,_,grey),
+    input_colour(R2,_,grey),
+    R < R2.
+
+bottom_wall(R):-
+    input_colour(R,_,grey),
+    input_colour(R2,_,grey),
+    R2 > R.
+
+left_wall_unique(W):-
+    setof(C, left_wall(C), L),
+    min(L, W).
+
+right_wall_unique(W):-
+    setof(C, right_wall(C), L),
+    min(L, W).
+
+top_wall_unique(W):-
+    setof(C, top_wall(C), L),
+    min(L, W).
+
+bottom_wall_unique(W):-
+    setof(C, bottom_wall(C), L),
+    min(L, W).
+
+
+gap_horizontal(R,C):-
+    input_colour(R,C,black),
+    RT is R + 1,
+    input_colour(RT,C,grey),
+    RB is R - 1,
+    input_colour(RB,C,grey).
+
+gap_vertical(R,C):-
+    input_colour(R,C,black),
+    CT is C + 1,
+    input_colour(R,CT,grey),
+    CB is C - 1,
+    input_colour(R,CB,grey).
+
+
+output_colour(R,C,grey):-
+    input_colour(R,C,grey).
+
+output_colour(R,C,teal):-
+    row(R),
+    column(C),
+    top_wall_unique(RT),
+    R < RT,
+    bottom_wall_unique(RB),
+    R > RB,
+    left_wall_unique(CT),
+    C > CT,
+    right_wall_unique(CB),
+    C < CB.
+
+output_colour(R,C,teal):-
+    row(R),
+    column(C),
+    gap_horizontal(R,CG),
+    CG =< C,
+    input_colour(R,CL,grey),
+    CL < C.
+
+output_colour(R,C,teal):-
+    row(R),
+    column(C),
+    gap_horizontal(R,CG),
+    CG >= C,
+    input_colour(R,CL,grey),
+    CL > C.
+
+output_colour(R,C,teal):-
+    row(R),
+    column(C),
+    gap_vertical(RG,C),
+    RG =< R,
+    input_colour(RB,C,grey),
+    RB < R.
+
+output_colour(R,C,teal):-
+    row(R),
+    column(C),
+    gap_vertical(RG,C),
+    RG >= R,
+    input_colour(RB,C,grey),
+    RB > R.

task.py

 import numpy as np
 
-from main import run_prolog_program, prolog2FOL_array, FOL2grid, grid2FOL, FOL2prolog
+from main import run_prolog_program, prolog2FOL_array, FOL2grid, grid2FOL, FOL2prolog, array_and_plot_grid
 
 
 class Grid:
@@ -28,6 +28,10 @@ class Example:
         result = same.all()
         if not result:
             print(f"output incorrect due to {(~same).sum()}/{np.prod(same.shape)} cells")
+            array_and_plot_grid(self.input_grid.grid)
+            array_and_plot_grid(out_grid)
+            array_and_plot_grid(self.output_grid.grid)
+            raise Exception
         return result
 
 class Task: