diff --git a/GALLERY.md b/GALLERY.md index be5a3a65..56e5836f 100644 --- a/GALLERY.md +++ b/GALLERY.md @@ -4,6 +4,7 @@ This gallery shows examples from all available datasets using their default conf ## Available Datasets - [advanced_geometry](#advanced_geometry) - [aiw](#aiw) +- [arc_1d](#arc_1d) - [base_conversion](#base_conversion) - [basic_arithmetic](#basic_arithmetic) - [bf](#bf) @@ -122,6 +123,88 @@ Metadata: {'task_type': 'friends'} ```` +### arc_1d +Generates ARC 1D tasks by randomly selecting from available task generators + +Default configuration: +```python +min_size = 10 +max_size = 30 +num_train = 3 +seed = 42 +size = 500 +``` + +Example tasks: +```` +Example 1: +Question: Find the common rule that maps an input grid to an output grid, given the examples below. + +Example 1: +Input: 7 1 0 0 5 5 0 5 5 0 0 0 0 +Output: 7 1 0 0 7 7 0 1 1 0 0 0 0 + +Example 2: +Input: 5 1 0 5 5 0 5 5 0 0 0 0 0 +Output: 5 1 0 5 5 0 1 1 0 0 0 0 0 + +Example 3: +Input: 2 6 0 0 5 5 0 5 5 0 0 0 0 +Output: 2 6 0 0 2 2 0 6 6 0 0 0 0 + +Below is a test input grid. Predict the corresponding output grid by applying the rule you found. Describe how you derived the rule and your overall reasoning process in detail before you submit your answer. Your final answer must be placed in tags and should be just be the text output grid itself. + +Input: +6 0 0 0 0 0 0 5 5 5 0 0 0 +Answer: 6 0 0 0 0 0 0 6 6 6 0 0 0 +Metadata: {'task_name': 'recolor_blocks_from_palette', 'size': 13, 'train_examples': [{'input': [7, 1, 0, 0, 5, 5, 0, 5, 5, 0, 0, 0, 0], 'output': [7, 1, 0, 0, 7, 7, 0, 1, 1, 0, 0, 0, 0]}, {'input': [5, 1, 0, 5, 5, 0, 5, 5, 0, 0, 0, 0, 0], 'output': [5, 1, 0, 5, 5, 0, 1, 1, 0, 0, 0, 0, 0]}, {'input': [2, 6, 0, 0, 5, 5, 0, 5, 5, 0, 0, 0, 0], 'output': [2, 6, 0, 0, 2, 2, 0, 6, 6, 0, 0, 0, 0]}], 'test_example': {'input': [6, 0, 0, 0, 0, 0, 0, 5, 5, 5, 0, 0, 0], 'output': [6, 0, 0, 0, 0, 0, 0, 6, 6, 6, 0, 0, 0]}} + +Example 2: +Question: Find the common rule that maps an input grid to an output grid, given the examples below. + +Example 1: +Input: 0 8 8 8 8 8 8 8 8 8 8 8 8 0 0 0 0 0 0 +Output: 0 0 0 0 8 8 8 8 8 8 8 8 8 8 8 8 0 0 0 + +Example 2: +Input: 0 0 0 0 0 0 0 0 0 0 2 2 2 2 2 2 0 0 0 +Output: 0 0 0 0 0 0 0 0 0 0 0 0 0 2 2 2 2 2 2 + +Example 3: +Input: 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 0 0 0 0 +Output: 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 0 + +Below is a test input grid. Predict the corresponding output grid by applying the rule you found. Describe how you derived the rule and your overall reasoning process in detail before you submit your answer. Your final answer must be placed in tags and should be just be the text output grid itself. + +Input: +0 0 0 0 0 0 6 6 6 6 6 6 6 6 6 0 0 0 0 +Answer: 0 0 0 0 0 0 0 0 0 6 6 6 6 6 6 6 6 6 0 +Metadata: {'task_name': 'move_3pix_solid', 'size': 19, 'train_examples': [{'input': [0, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 0, 0, 0, 0, 0, 0], 'output': [0, 0, 0, 0, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 0, 0, 0]}, {'input': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 2, 2, 2, 2, 2, 0, 0, 0], 'output': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 2, 2, 2, 2, 2]}, {'input': [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0], 'output': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0]}], 'test_example': {'input': [0, 0, 0, 0, 0, 0, 6, 6, 6, 6, 6, 6, 6, 6, 6, 0, 0, 0, 0], 'output': [0, 0, 0, 0, 0, 0, 0, 0, 0, 6, 6, 6, 6, 6, 6, 6, 6, 6, 0]}} + +Example 3: +Question: Find the common rule that maps an input grid to an output grid, given the examples below. + +Example 1: +Input: 0 0 0 0 0 0 0 2 0 0 4 4 4 4 4 4 4 4 4 4 4 4 4 0 0 0 +Output: 0 0 0 0 0 0 0 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 0 0 0 + +Example 2: +Input: 0 0 0 2 0 0 0 0 0 0 0 0 0 3 3 3 3 3 3 3 3 0 0 0 0 0 +Output: 0 0 0 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 0 0 0 0 0 + +Example 3: +Input: 0 0 0 0 0 2 0 0 0 0 0 0 0 0 0 0 0 0 3 3 3 3 0 0 0 0 +Output: 0 0 0 0 0 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 0 0 0 0 + +Below is a test input grid. Predict the corresponding output grid by applying the rule you found. Describe how you derived the rule and your overall reasoning process in detail before you submit your answer. Your final answer must be placed in tags and should be just be the text output grid itself. + +Input: +0 0 0 0 0 0 0 0 0 0 0 7 7 7 7 7 7 7 7 7 7 7 7 7 2 0 +Answer: 0 0 0 0 0 0 0 0 0 0 0 7 7 7 7 7 7 7 7 7 7 7 7 7 7 0 +Metadata: {'task_name': 'block_scale_to_dot', 'size': 26, 'train_examples': [{'input': [0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 0, 0, 0], 'output': [0, 0, 0, 0, 0, 0, 0, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 0, 0, 0]}, {'input': [0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 3, 3, 3, 3, 3, 3, 3, 0, 0, 0, 0, 0], 'output': [0, 0, 0, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 0, 0, 0, 0, 0]}, {'input': [0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 3, 3, 3, 0, 0, 0, 0], 'output': [0, 0, 0, 0, 0, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 0, 0, 0, 0]}], 'test_example': {'input': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 2, 0], 'output': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 0]}} + +```` + ### base_conversion Generates base conversion tasks @@ -1368,7 +1451,7 @@ Example tasks: ```` Example 1: Question: In front of you are some buttons, a light, and a number. The light will toggle between red and green whenever you press a button. Each button performs a mathematical operation to the number, but the operation may depend on the state of the light. -You must press the shortest correct sequence of buttons to reach the target value. +You must press the shortest correct sequence of buttons to reach the target value. Your answer should be a sequence of buttons separated by '→', for example: A → B → C Start: 0 (red) Target: 46 @@ -1381,7 +1464,7 @@ Metadata: {'difficulty': 10, 'solution_path': ['A', 'B', 'C', 'C', 'A', 'C'], 't Example 2: Question: In front of you are some buttons, a light, and a number. The light will toggle between red and green whenever you press a button. Each button performs a mathematical operation to the number, but the operation may depend on the state of the light. -You must press the shortest correct sequence of buttons to reach the target value. +You must press the shortest correct sequence of buttons to reach the target value. Your answer should be a sequence of buttons separated by '→', for example: A → B → C Start: 0 (red) Target: 30 @@ -1394,7 +1477,7 @@ Metadata: {'difficulty': 10, 'solution_path': ['C', 'A', 'C', 'A', 'C', 'A', 'C' Example 3: Question: In front of you are some buttons, a light, and a number. The light will toggle between red and green whenever you press a button. Each button performs a mathematical operation to the number, but the operation may depend on the state of the light. -You must press the shortest correct sequence of buttons to reach the target value. +You must press the shortest correct sequence of buttons to reach the target value. Your answer should be a sequence of buttons separated by '→', for example: A → B → C Start: 0 (red) Target: 45 diff --git a/reasoning_gym/__init__.py b/reasoning_gym/__init__.py index 49068294..054cbd95 100644 --- a/reasoning_gym/__init__.py +++ b/reasoning_gym/__init__.py @@ -5,7 +5,7 @@ from . import algebra, algorithmic, arithmetic, cognition, data, games, geometry, graphs, logic from .factory import create_dataset, register_dataset -__version__ = "0.1.1" +__version__ = "0.1.3" __all__ = [ "algebra", "algorithmic", diff --git a/reasoning_gym/cognition/__init__.py b/reasoning_gym/cognition/__init__.py index fddd97b1..38baf31b 100644 --- a/reasoning_gym/cognition/__init__.py +++ b/reasoning_gym/cognition/__init__.py @@ -6,18 +6,21 @@ - Working memory """ +from .arc_1d import Arc1DConfig, Arc1DDataset from .color_cube_rotation import ColorCubeRotationConfig, ColorCubeRotationDataset from .figlet_fonts import FigletFontConfig, FigletFontDataset from .number_sequences import NumberSequenceConfig, NumberSequenceDataset from .rubiks_cube import RubiksCubeConfig, RubiksCubeDataset __all__ = [ - "NumberSequenceConfig", - "NumberSequenceDataset", + "Arc1DConfig", + "Arc1DDataset", "ColorCubeRotationConfig", "ColorCubeRotationDataset", - "RubiksCubeConfig", - "RubiksCubeDataset", "FigletFontConfig", "FigletFontDataset", + "NumberSequenceConfig", + "NumberSequenceDataset", + "RubiksCubeConfig", + "RubiksCubeDataset", ] diff --git a/reasoning_gym/cognition/arc_1d.py b/reasoning_gym/cognition/arc_1d.py index badea2c4..332a5851 100644 --- a/reasoning_gym/cognition/arc_1d.py +++ b/reasoning_gym/cognition/arc_1d.py @@ -1,1024 +1,17 @@ -import random +from dataclasses import dataclass from random import Random -from typing import Dict, List, Optional, Callable, Tuple +from typing import Optional from ..dataset import ProceduralDataset from ..factory import register_dataset -def gen_field(size: int, color: int = 0) -> List[int]: - """Generate a field of given size filled with specified color (default 0).""" - return [color] * size - - -def write_block(pos: int, block: List[int], field: List[int]) -> List[int]: - """Write a block into a field at given position.""" - result = field.copy() - for i, color in enumerate(block): - result[pos + i] = color - return result - - -def task_move_n_pix(rng: Random, size: int, move_pix: int, solid: bool) -> Optional[Dict[str, List[int]]]: - """Generate a task where a block is moved to the right by move_pix pixels.""" - if size <= move_pix + 1: - return None - - block_size = rng.randint(1, size - move_pix - 1) - block_pos = rng.randint(0, size - block_size - move_pix) - - if solid: - color = rng.randint(1, 9) - block = [color] * block_size - else: - block = [rng.randint(1, 9) for _ in range(block_size)] - - question = write_block(block_pos, block, gen_field(size)) - answer = write_block(block_pos + move_pix, block, gen_field(size)) - - return {"input": question, "output": answer} - - -def task_move_n_pix_wrapped(rng: Random, size: int, move_pix: int, solid: bool) -> Optional[Dict[str, List[int]]]: - """Generate a task where a block is moved to the right by move_pix pixels with wrapping.""" - block_size = rng.randint(1, size) - block_pos = rng.randint(0, size) - - if solid: - color = rng.randint(1, 9) - block = [color] * block_size - else: - block = [rng.randint(1, 9) for _ in range(block_size)] - - question = gen_field(size) - answer = gen_field(size) - - for i, color in enumerate(block): - question[(block_pos + i) % size] = color - answer[(block_pos + move_pix + i) % size] = color - - return {"input": question, "output": answer} - - -def task_gravity(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where all non-zero elements are attracted to the left.""" - density = 0.5 - question = [rng.randint(1, 9) if rng.random() < density else 0 for _ in range(size)] - - non_zero = [x for x in question if x != 0] - answer = non_zero + [0] * (size - len(non_zero)) - - return {"input": question, "output": answer} - - -def task_gravity_counting(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where non-zero elements are counted and represented as a sequence of 1s.""" - density = 0.5 - question = [rng.randint(1, 9) if rng.random() < density else 0 for _ in range(size)] - - count = sum(1 for x in question if x != 0) - answer = [1] * count + [0] * (size - count) - - return {"input": question, "output": answer} - - -def task_gravity_antigravity(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where color 1 moves right and color 2 moves left.""" - density = 0.5 - question = [rng.randint(1, 2) if rng.random() < density else 0 for _ in range(size)] - - color1 = [x for x in question if x == 1] - color2 = [x for x in question if x == 2] - answer = [2] * len(color2) + [0] * (size - len(color1) - len(color2)) + [1] * len(color1) - - return {"input": question, "output": answer} - - -def task_block_touch_dot(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where a block moves to touch (but not cover) a dot.""" - dot_color = 1 - block_color = rng.randint(2, 9) - - block_size = rng.randint(1, size) - dot_pos = rng.randint(0, size) - - can_place_left = dot_pos >= block_size - can_place_right = dot_pos + block_size < size - - if not (can_place_left or can_place_right): - return None - - if can_place_left and can_place_right: - side = rng.choice(["left", "right"]) - elif can_place_left: - side = "left" - else: - side = "right" - - if side == "left": - q_block_pos = rng.randint(0, dot_pos - block_size) - a_block_pos = dot_pos - block_size - else: - q_block_pos = rng.randint(dot_pos + 1, size - block_size) - a_block_pos = dot_pos + 1 - - question = gen_field(size) - question[dot_pos] = dot_color - question = write_block(q_block_pos, [block_color] * block_size, question) - - answer = gen_field(size) - answer[dot_pos] = dot_color - answer = write_block(a_block_pos, [block_color] * block_size, answer) - - return {"input": question, "output": answer} - - -def task_block_touch_dot_n_pix(rng: Random, size: int, move_pix: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where a block moves move_pix pixels toward a dot.""" - dot_color = 2 - block_color = rng.randint(3, 9) - - block_size = rng.randint(1, size) - dot_pos = rng.randint(0, size) - - can_place_left = dot_pos >= block_size - can_place_right = dot_pos + block_size < size - - if not (can_place_left or can_place_right): - return None - - if can_place_left and can_place_right: - side = rng.choice(["left", "right"]) - elif can_place_left: - side = "left" - else: - side = "right" - - if side == "left": - q_block_pos = rng.randint(0, dot_pos - block_size) - distance = (dot_pos - block_size) - q_block_pos - move = min(distance, move_pix) - a_block_pos = q_block_pos + move - else: - q_block_pos = rng.randint(dot_pos + 1, size - block_size) - distance = q_block_pos - (dot_pos + 1) - move = min(distance, move_pix) - a_block_pos = q_block_pos - move - - question = gen_field(size) - question[dot_pos] = dot_color - question = write_block(q_block_pos, [block_color] * block_size, question) - - answer = gen_field(size) - answer[dot_pos] = dot_color - answer = write_block(a_block_pos, [block_color] * block_size, answer) - - return {"input": question, "output": answer} - - -def task_block_scale_to_dot(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where a block scales to touch a dot (keeping one end fixed).""" - dot_color = 2 - block_color = rng.randint(3, 9) - - block_size = rng.randint(1, size) - dot_pos = rng.randint(0, size) - - can_place_left = dot_pos >= block_size - can_place_right = dot_pos + block_size < size - - if not (can_place_left or can_place_right): - return None - - if can_place_left and can_place_right: - side = rng.choice(["left", "right"]) - elif can_place_left: - side = "left" - else: - side = "right" - - if side == "left": - q_block_pos = rng.randint(0, dot_pos - block_size) - new_size = dot_pos - q_block_pos + 1 - a_block_pos = q_block_pos - else: - q_block_pos = rng.randint(dot_pos + 1, size - block_size) - new_size = (q_block_pos + block_size) - dot_pos - a_block_pos = dot_pos - - question = gen_field(size) - question[dot_pos] = dot_color - question = write_block(q_block_pos, [block_color] * block_size, question) - - answer = gen_field(size) - answer[dot_pos] = dot_color - answer = write_block(a_block_pos, [block_color] * new_size, answer) - - return {"input": question, "output": answer} - - -def task_two_points_and_fill(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where space between two points of same color is filled with that color.""" - color = rng.randint(1, 9) - - pos1 = rng.randint(0, size - 1) - pos2 = rng.randint(0, size - 1) - if pos1 == pos2: - return None - - pos1, pos2 = min(pos1, pos2), max(pos1, pos2) - - question = gen_field(size) - question[pos1] = color - question[pos2] = color - - answer = question.copy() - for i in range(pos1, pos2 + 1): - answer[i] = color - - return {"input": question, "output": answer} - - -def task_reflect_block_with_border_pixel(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where a block with a border pixel is reflected.""" - block_size = rng.randint(2, size) - if block_size > size: - return None - - c1 = rng.randint(1, 9) - c2 = rng.randint(1, 9) - if c1 == c2: - return None - - side = "left" if rng.random() < 0.5 else "right" - pos = rng.randint(0, size - block_size) - - block = [c1] * block_size - if side == "left": - block[0] = c2 - else: - block[block_size - 1] = c2 - - question = write_block(pos, block, gen_field(size)) - reversed_block = block[::-1] # Reverse the block - answer = write_block(pos, reversed_block, gen_field(size)) - - return {"input": question, "output": answer} - - -def task_reflect_block_with_border_pixel_random(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where a random-colored block with a border pixel is reflected.""" - block_size = rng.randint(2, size) - if block_size > size: - return None - - side = "left" if rng.random() < 0.5 else "right" - pos = rng.randint(0, size - block_size) - - block = [rng.randint(1, 9) for _ in range(block_size)] - border_color = rng.randint(1, 9) - - if side == "left": - if block[0] == border_color: - return None - block[0] = border_color - else: - if block[block_size - 1] == border_color: - return None - block[block_size - 1] = border_color - - question = write_block(pos, block, gen_field(size)) - reversed_block = block[::-1] # Reverse the block - answer = write_block(pos, reversed_block, gen_field(size)) - - return {"input": question, "output": answer} - - -def task_reflect_block_around_dot(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where a block is reflected around a dot.""" - dot_color = 2 - - dot_pos = rng.randint(0, size) - block_size = rng.randint(1, size) - block_pos = rng.randint(0, size - block_size) - block_end = block_pos + block_size - 1 - - # Check if block is strictly to left or right of dot - strictly_left = block_end < dot_pos - strictly_right = block_pos > dot_pos - - if not (strictly_left or strictly_right): - return None - - block_color = rng.randint(3, 9) # Different from dot color - block = [block_color] * block_size - - # Calculate reflection bounds - min_reflect = 2 * dot_pos - block_end - max_reflect = 2 * dot_pos - block_pos - if min_reflect < 0 or max_reflect >= size: - return None - - question = gen_field(size) - question = write_block(block_pos, block, question) - question[dot_pos] = dot_color - - answer = gen_field(size) - answer[dot_pos] = dot_color - for i in range(block_size): - reflect_idx = 2 * dot_pos - (block_pos + i) - answer[reflect_idx] = block[i] - - return {"input": question, "output": answer} - - -def task_block_and_noise_remove(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where noise around a block needs to be removed.""" - block_size = rng.randint(2, size) - if block_size > size: - return None - - block_pos = rng.randint(0, size - block_size) - color = rng.randint(1, 9) - - # Create field with block - field = gen_field(size) - for i in range(block_size): - field[block_pos + i] = color - - # Track forbidden positions for noise - forbidden = [False] * size - for i in range(block_pos, block_pos + block_size): - forbidden[i] = True - if block_pos > 0: - forbidden[block_pos - 1] = True - if block_pos + block_size < size: - forbidden[block_pos + block_size] = True - - # Add noise - noise_count = rng.randint(1, 3) - noise_positions = [] - - for _ in range(noise_count): - allowed = [i for i in range(size) if not forbidden[i]] - if not allowed: - break - noise_pos = rng.choice(allowed) - noise_positions.append(noise_pos) - field[noise_pos] = color - forbidden[noise_pos] = True - if noise_pos > 0: - forbidden[noise_pos - 1] = True - if noise_pos + 1 < size: - forbidden[noise_pos + 1] = True - - if len(noise_positions) < noise_count: - return None - - question = field - answer = field.copy() - for pos in noise_positions: - answer[pos] = 0 - - return {"input": question, "output": answer} - - -def task_block_and_noise_remove_inside(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where noise inside a block needs to be removed.""" - if size <= 6: - return None - - block_size = rng.randint(6, size) - if block_size > size: - return None - - block_pos = rng.randint(0, size - block_size) - color = rng.randint(1, 9) - - # Create field with block - field = gen_field(size) - for i in range(block_size): - field[block_pos + i] = color - - # Add noise inside block - max_noise = max(1, (block_size // 2) - 1) - noise_count = rng.randint(1, max_noise) - - positions = list(range(block_size)) - rng.shuffle(positions) - noise_positions = positions[:noise_count] - - for offset in noise_positions: - pos = block_pos + offset - noise_color = rng.randint(1, 9) - while noise_color == color: - noise_color = rng.randint(1, 9) - field[pos] = noise_color - - question = field - answer = field.copy() - for offset in noise_positions: - answer[block_pos + offset] = color - - return {"input": question, "output": answer} - - -def task_copy_block_to_dots(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where a block pattern is copied to dot positions.""" - block_size = 3 if rng.random() < 0.5 else 5 - if block_size >= size: - return None - - color = rng.randint(1, 9) - block = [color] * block_size - - # Generate dots with minimum distance to prevent overlap - min_gap = block_size - dot_positions = [] - pos = block_size + block_size // 2 + 1 - - while pos <= size - block_size: - if rng.random() < 0.5: # Control dot density - dot_positions.append(pos) - pos += min_gap - pos += 1 - - if not dot_positions: - return None - - question = gen_field(size) - question = write_block(0, block, question) - for pos in dot_positions: - question[pos] = color - - answer = gen_field(size) - answer = write_block(0, block, answer) - for pos in dot_positions: - block_start = pos - block_size // 2 - answer = write_block(block_start, block, answer) - - return {"input": question, "output": answer} - - -def task_copy_block_to_dots_colors(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where a block pattern is copied to dot positions with matching colors.""" - block_size = 3 if rng.random() < 0.5 else 5 - if block_size >= size: - return None - - block_color = rng.randint(1, 9) - block = [block_color] * block_size - - # Generate dots with minimum distance to prevent overlap - min_gap = block_size - dot_positions = [] - dot_colors = [] - pos = block_size + block_size // 2 + 1 - - while pos < size - block_size: - if rng.random() < 0.5: - dot_color = rng.randint(1, 9) - dot_positions.append(pos) - dot_colors.append(dot_color) - pos += min_gap - pos += 1 - - if not dot_positions: - return None - - question = gen_field(size) - question = write_block(0, block, question) - for i, pos in enumerate(dot_positions): - question[pos] = dot_colors[i] - - answer = gen_field(size) - answer = write_block(0, block, answer) - for i, pos in enumerate(dot_positions): - block_start = pos - block_size // 2 - colored_block = [dot_colors[i]] * block_size - answer = write_block(block_start, colored_block, answer) - - return {"input": question, "output": answer} - - -def task_paint_biggest_block(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where the largest block is painted a different color.""" - target_color = 1 - initial_color = rng.randint(2, 9) - - # Generate random blocks - question = gen_field(size) - blocks = [] - pos = 0 - - while pos < size: - if rng.random() < 0.4 and size - pos >= 2: - block_size = rng.randint(2, min(size - pos, 6)) - blocks.append((pos, block_size)) - for i in range(block_size): - question[pos + i] = initial_color - pos += block_size + 1 - else: - pos += 1 - - if len(blocks) < 2: - return None - - # Find biggest block - biggest_pos, biggest_size = max(blocks, key=lambda x: x[1]) - - # Check if there are multiple blocks of the same size - biggest_count = sum(1 for _, size in blocks if size == biggest_size) - if biggest_count > 1: - return None - - answer = question.copy() - for i in range(biggest_size): - answer[biggest_pos + i] = target_color - - return {"input": question, "output": answer} - - -def task_sort_blocks_by_size(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where blocks are sorted by size with 1 pixel gaps.""" - color = rng.randint(1, 9) - blocks = [] - pos = 0 - - # Generate random blocks with random sizes - while pos < size: - if rng.random() < 0.4 and size - pos >= 2: - block_size = rng.randint(1, min(size - pos, 6)) - blocks.append((pos, block_size)) - pos += block_size + rng.randint(1, 4) # Random gaps - else: - pos += 1 - - if len(blocks) < 2: - return None - - # Create input field - question = gen_field(size) - for pos, block_size in blocks: - for i in range(block_size): - question[pos + i] = color - - # Sort blocks by size - blocks.sort(key=lambda x: x[1]) - - # Check if sorted blocks fit with gaps - total_space = sum(size for _, size in blocks) + len(blocks) - 1 - if total_space > size: - return None - - # Create answer field with sorted blocks - answer = gen_field(size) - current_pos = 0 - - for _, block_size in blocks: - for i in range(block_size): - answer[current_pos + i] = color - current_pos += block_size + 1 # One pixel gap - - return {"input": question, "output": answer} - - -def task_sort_complete_sequence(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where a complete sequence of block sizes is sorted.""" - # Calculate max possible block size given total array size - max_size = 1 - total_space = 0 - while total_space + max_size + 1 <= size: - total_space += max_size + 1 - max_size += 1 - max_size -= 1 - - if max_size < 2: - return None - - color = rng.randint(1, 9) - - # Create sequence of all sizes from 1 to max_size - blocks = list(range(1, max_size + 1)) - rng.shuffle(blocks) - - # Create input field with shuffled blocks - question = gen_field(size) - pos = 0 - for block_size in blocks: - for i in range(block_size): - question[pos + i] = color - pos += block_size + 1 - - # Create answer field with sorted blocks - answer = gen_field(size) - pos = 0 - for block_size in range(1, max_size + 1): - for i in range(block_size): - answer[pos + i] = color - pos += block_size + 1 - - return {"input": question, "output": answer} - - -def task_recolor_blocks_by_size(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where two blocks are recolored based on their size.""" - # Generate two different random sizes - size1 = rng.randint(2, 8) - size2 = rng.randint(2, 8) - while size2 == size1: - size2 = rng.randint(2, 8) - - # Ensure both blocks fit with at least 1 gap - if size1 + size2 + 1 > size: - return None - - # Place blocks with gap - pos1 = rng.randint(0, size - (size1 + size2 + 1)) - pos2 = rng.randint(pos1 + size1 + 1, size - size2) - - # Create input field with both blocks color 3 - question = gen_field(size) - for i in range(size1): - question[pos1 + i] = 3 - for i in range(size2): - question[pos2 + i] = 3 - - # Create answer field with recolored blocks - answer = question.copy() - if size1 > size2: - for i in range(size1): - answer[pos1 + i] = 1 - for i in range(size2): - answer[pos2 + i] = 2 - else: - for i in range(size1): - answer[pos1 + i] = 2 - for i in range(size2): - answer[pos2 + i] = 1 - - return {"input": question, "output": answer} - - -def task_gravity_one_step(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where non-zero elements move one step left if possible.""" - question = [rng.randint(1, 9) if rng.random() < 0.5 else 0 for _ in range(size)] - answer = question.copy() - - # Move each non-zero pixel one step left if possible - for i in range(1, size): - if answer[i] != 0 and answer[i - 1] == 0: - answer[i - 1] = answer[i] - answer[i] = 0 - - return {"input": question, "output": answer} - - -def task_move_block_by_own_size(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where a block moves right by its own size.""" - block_size = rng.randint(1, size // 2) # Ensure space for movement - pos = rng.randint(0, size - block_size * 2) # Space for block and movement - color = rng.randint(1, 9) - - question = gen_field(size) - block = [color] * block_size - question = write_block(pos, block, question) - - answer = write_block(pos + block_size, block, gen_field(size)) - - return {"input": question, "output": answer} - - -def task_change_to_five(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where all non-zero colors change to 5.""" - density = 0.5 - question = [rng.randint(1, 9) if rng.random() < density else 0 for _ in range(size)] - answer = [5 if x != 0 else 0 for x in question] - - return {"input": question, "output": answer} - - -def task_recolor_blocks_from_palette(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where blocks are recolored using a color palette.""" - # Generate blocks of same size - block_size = rng.randint(2, 4) - blocks = [] - pos = 0 - - while pos + block_size <= size: - if rng.random() < 0.4: - blocks.append(pos) - pos += block_size + 1 - else: - pos += 1 - - # Ensure we have space for palette - while blocks and blocks[-1] + block_size + len(blocks) + 1 >= size: - blocks.pop() - - if not blocks: - return None - - # Shift blocks right to make room for palette - palette_size = len(blocks) - blocks = [pos + palette_size + 1 for pos in blocks] - - # Generate color palette - colors = [] - for _ in range(len(blocks)): - while True: - color = rng.randint(1, 9) - if color not in colors: - colors.append(color) - break - - # Create question with color palette and blocks - question = gen_field(size) - - # Place color palette at start - for i, color in enumerate(colors): - question[i] = color - - # Place blocks of color 5 - for block_pos in blocks: - for i in range(block_size): - question[block_pos + i] = 5 - - # Create answer with recolored blocks - answer = question.copy() - for block_idx, block_pos in enumerate(blocks): - color = colors[block_idx] - for i in range(block_size): - answer[block_pos + i] = color - - return {"input": question, "output": answer} - - -def task_duplicate_block_from_seeds(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where a block is duplicated from seed pixels.""" - block_size = rng.randint(2, 4) - if block_size + 1 >= size: - return None - if size <= 3 + block_size: - return None - - # Position block with space for seeds - block_pos = rng.randint(2, size - block_size - 1) - - # Decide seed placement - left_seed = rng.random() < 0.5 - right_seed = rng.random() < 0.5 - if not (left_seed or right_seed): - return None - - # Create input - question = gen_field(size) - - # Place main block - for i in range(block_size): - question[block_pos + i] = 1 - - # Place seeds with gaps - seeds = [] - if left_seed: - color = rng.randint(1, 9) - question[block_pos - 2] = color - seeds.append(("left", block_pos - 2, color)) - if right_seed: - color = rng.randint(1, 9) - question[block_pos + block_size + 1] = color - seeds.append(("right", block_pos + block_size + 1, color)) - - # Create answer with duplicated blocks - answer = question.copy() - - for side, seed_pos, color in seeds: - if side == "left": - # For left seed, blocks end at seed - end_pos = seed_pos - while end_pos >= 0: - start_pos = end_pos - block_size + 1 - for pos in range(max(0, start_pos), end_pos + 1): - answer[pos] = color - if start_pos < 1: - break - end_pos = start_pos - 2 # -1 for gap - else: # side == "right" - # For right seed, blocks start at seed - start_pos = seed_pos - while start_pos < size: - for offset in range(min(block_size, size - start_pos)): - answer[start_pos + offset] = color - if start_pos + block_size + 1 >= size: - break - start_pos = start_pos + block_size + 1 # +1 for gap - - return {"input": question, "output": answer} - - -def task_fill_from_pixel(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where a pixel fills in one direction until hitting another pixel.""" - block_size = rng.randint(3, 6) - if block_size >= size - 2: - return None - - # Position block with space for seed - block_pos = rng.randint(1, size - block_size - 1) - - # Create input - question = gen_field(size) - - # Place main block - block_color = rng.randint(1, 9) - for i in range(block_size): - question[block_pos + i] = block_color - - # Place seed pixel and determine fill direction - seed_color = rng.randint(1, 9) - while seed_color == block_color: - seed_color = rng.randint(1, 9) - - is_left = rng.random() < 0.5 - - if is_left: - question[block_pos - 1] = seed_color - else: - question[block_pos + block_size] = seed_color - - # Create answer with fill - answer = question.copy() - - if is_left: - # Fill from seed to left border - for i in range(block_pos): - answer[i] = seed_color - else: - # Fill from seed to right border - for i in range(block_pos + block_size, size): - answer[i] = seed_color - - return {"input": question, "output": answer} - - -def task_mark_size_two_blocks(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where size-2 blocks are marked with surrounding pixels.""" - blocks = [] - pos = 0 - - # Generate blocks with minimum gap of 2 - while pos < size: - if rng.random() < 0.4: - block_size = rng.randint(1, 3) - # Check if we have space for block and potential markers - needed_space = block_size + (2 if block_size == 2 else 0) - if pos + needed_space < size: - blocks.append((pos, block_size)) - pos += block_size + 2 # Minimum gap of 2 - - pos += 1 - - if len(blocks) < 2: - return None - - # Verify gaps between blocks (including markers) - valid = True - for i in range(len(blocks) - 1): - pos1, size1 = blocks[i] - pos2, _ = blocks[i + 1] - needed_gap = 3 if size1 == 2 else 2 - if pos2 - (pos1 + size1) < needed_gap: - valid = False - break - if not valid: - return None - - # Create input with blocks - question = gen_field(size) - for pos, block_size in blocks: - # Place block - for i in range(block_size): - question[pos + i] = 1 - - # Create answer with markers - answer = question.copy() - for pos, block_size in blocks: - if block_size == 2: - # Add markers for size 2 blocks - if pos > 0: - answer[pos - 1] = 3 - if pos + block_size < size: - answer[pos + block_size] = 3 - - return {"input": question, "output": answer} - - -def task_fill_until_collision(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where pixels fill empty space until collision.""" - # At least 4 positions for meaningful puzzle - if size < 4: - return None - - is_left = rng.random() < 0.5 - question = gen_field(size) - - # Place the side marker - if is_left: - question[0] = 5 - else: - question[size - 1] = 5 - - # Place 2-4 random pixels - num_pixels = rng.randint(2, 4) - positions = [] - - if is_left: - # Skip first position - for _ in range(num_pixels): - while True: - pos = rng.randint(1, size - 1) - if pos not in positions: - positions.append(pos) - break - else: - # Skip last position - for _ in range(num_pixels): - while True: - pos = rng.randint(0, size - 2) - if pos not in positions: - positions.append(pos) - break - - # Color random pixels - for pos in positions: - question[pos] = rng.randint(1, 9) - - positions.sort() - - # Create answer - answer = question.copy() - - if is_left: - # Fill right from each pixel - prev_pos = 0 # Start from marker - for pos in positions: - color = question[pos] - # Fill from previous position to current - for i in range(prev_pos + 1, pos): - answer[i] = color - prev_pos = pos - else: - # Fill left from each pixel - prev_pos = size - 1 # Start from marker - for pos in reversed(positions): - color = question[pos] - # Fill from current position to previous - for i in range(pos + 1, prev_pos): - answer[i] = color - prev_pos = pos - - return {"input": question, "output": answer} - - -def task_repeat_pattern_full(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where a pattern is repeated to fill the space.""" - # Generate initial pattern - pattern_size = rng.randint(2, 5) - pattern = [rng.randint(1, 9) for _ in range(pattern_size)] - - # Calculate total size needed for 2 repetitions - double_size = pattern_size * 2 - if double_size >= size: - return None - - # Create input with 2 repetitions - question = gen_field(size) - for i in range(pattern_size): - question[i] = pattern[i] - question[i + pattern_size] = pattern[i] - - # Create answer with maximum repetitions - answer = gen_field(size) - pos = 0 - while pos + pattern_size <= size: - for i in range(pattern_size): - answer[pos + i] = pattern[i] - pos += pattern_size - - # Fill remaining space (if any) with pattern elements - for i in range(pos, size): - answer[i] = pattern[i - pos] - - return {"input": question, "output": answer} - - -from dataclasses import dataclass -from typing import Callable, Dict, List, Optional, Tuple - @dataclass class Arc1DConfig: """Configuration for ARC 1D task generation""" + min_size: int = 10 # Minimum grid size - max_size: int = 30 # Maximum grid size + max_size: int = 30 # Maximum grid size num_train: int = 3 # Number of training examples seed: Optional[int] = None size: int = 500 @@ -1031,139 +24,14 @@ def validate(self) -> None: assert self.size > 0, "size must be positive" -def task_gravity_weighted_colors(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where color 2 is heavier than color 1 in gravity.""" - # Generate random field with only colors 1 and 2 - question = [rng.randint(1, 2) if rng.random() < 0.5 else 0 for _ in range(size)] - - # Count colors - count_1 = sum(1 for x in question if x == 1) - count_2 = sum(1 for x in question if x == 2) - - # Create answer with sorted colors - answer = gen_field(size) - - # Place heavier color 2 first - for i in range(count_2): - answer[i] = 2 - - # Then place color 1 - for i in range(count_1): - answer[count_2 + i] = 1 - - return {"input": question, "output": answer} - - -def task_color_left_half_blocks(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: - """Generate a task where left half of blocks are colored differently.""" - pos = 0 - question = gen_field(size) - blocks = [] - - # Generate blocks with gap 1 - while pos < size: - if rng.random() < 0.4: - block_size = rng.randint(2, 8) - if pos + block_size >= size: - break - - blocks.append((pos, block_size)) - for i in range(block_size): - question[pos + i] = 2 - pos += block_size + 1 # block size + gap - else: - pos += 1 - - if len(blocks) < 2: - return None - - # Create answer with half-colored blocks - answer = question.copy() - for pos, block_size in blocks: - half_size = block_size // 2 - for i in range(half_size): - answer[pos + i] = 8 - - return {"input": question, "output": answer} - - -# Table of all ARC 1D task functions with their parameters -ARC_1D_TASKS = { - # Move tasks - "move_1pix_solid": (task_move_n_pix, {"move_pix": 1, "solid": True}), - "move_2pix_solid": (task_move_n_pix, {"move_pix": 2, "solid": True}), - "move_3pix_solid": (task_move_n_pix, {"move_pix": 3, "solid": True}), - "move_4pix_solid": (task_move_n_pix, {"move_pix": 4, "solid": True}), - "move_1pix_colorful": (task_move_n_pix, {"move_pix": 1, "solid": False}), - "move_2pix_colorful": (task_move_n_pix, {"move_pix": 2, "solid": False}), - "move_3pix_colorful": (task_move_n_pix, {"move_pix": 3, "solid": False}), - "move_4pix_colorful": (task_move_n_pix, {"move_pix": 4, "solid": False}), - - # Move wrapped tasks - "move_1pix_solid_wrapped": (task_move_n_pix_wrapped, {"move_pix": 1, "solid": True}), - "move_2pix_solid_wrapped": (task_move_n_pix_wrapped, {"move_pix": 2, "solid": True}), - "move_3pix_solid_wrapped": (task_move_n_pix_wrapped, {"move_pix": 3, "solid": True}), - "move_4pix_solid_wrapped": (task_move_n_pix_wrapped, {"move_pix": 4, "solid": True}), - "move_1pix_colorful_wrapped": (task_move_n_pix_wrapped, {"move_pix": 1, "solid": False}), - "move_2pix_colorful_wrapped": (task_move_n_pix_wrapped, {"move_pix": 2, "solid": False}), - "move_3pix_colorful_wrapped": (task_move_n_pix_wrapped, {"move_pix": 3, "solid": False}), - "move_4pix_colorful_wrapped": (task_move_n_pix_wrapped, {"move_pix": 4, "solid": False}), - - # Gravity tasks - "gravity": (task_gravity, {}), - "gravity_counting": (task_gravity_counting, {}), - "gravity_antigravity": (task_gravity_antigravity, {}), - "gravity_one_step": (task_gravity_one_step, {}), - "gravity_weighted_colors": (task_gravity_weighted_colors, {}), - - # Block tasks - "block_touch_dot": (task_block_touch_dot, {}), - "block_touch_dot_1pix": (task_block_touch_dot_n_pix, {"move_pix": 1}), - "block_touch_dot_2pix": (task_block_touch_dot_n_pix, {"move_pix": 2}), - "block_touch_dot_3pix": (task_block_touch_dot_n_pix, {"move_pix": 3}), - "block_touch_dot_4pix": (task_block_touch_dot_n_pix, {"move_pix": 4}), - "block_scale_to_dot": (task_block_scale_to_dot, {}), - "block_and_noise_remove": (task_block_and_noise_remove, {}), - "block_and_noise_remove_inside": (task_block_and_noise_remove_inside, {}), - "move_block_by_own_size": (task_move_block_by_own_size, {}), - - # Pattern tasks - "two_points_and_fill": (task_two_points_and_fill, {}), - "copy_block_to_dots": (task_copy_block_to_dots, {}), - "copy_block_to_dots_colors": (task_copy_block_to_dots_colors, {}), - "repeat_pattern_full": (task_repeat_pattern_full, {}), - - # Reflection tasks - "reflect_block_with_border_pixel": (task_reflect_block_with_border_pixel, {}), - "reflect_block_random": (task_reflect_block_with_border_pixel_random, {}), - "reflect_block_around_dot": (task_reflect_block_around_dot, {}), - - # Color tasks - "paint_biggest_block": (task_paint_biggest_block, {}), - "recolor_blocks_by_size": (task_recolor_blocks_by_size, {}), - "change_to_five": (task_change_to_five, {}), - "recolor_blocks_from_palette": (task_recolor_blocks_from_palette, {}), - "color_left_half_blocks": (task_color_left_half_blocks, {}), - - # Sorting tasks - "sort_blocks_by_size": (task_sort_blocks_by_size, {}), - "sort_complete_sequence": (task_sort_complete_sequence, {}), - - # Fill tasks - "duplicate_block_from_seeds": (task_duplicate_block_from_seeds, {}), - "fill_from_pixel": (task_fill_from_pixel, {}), - "fill_until_collision": (task_fill_until_collision, {}), - - # Marking tasks - "mark_size_two_blocks": (task_mark_size_two_blocks, {}), -} - - class Arc1DDataset(ProceduralDataset): """Generates ARC 1D tasks by randomly selecting from available task generators""" def __init__(self, config: Arc1DConfig): + from .arc_1d_tasks import ARC_1D_TASKS + super().__init__(config=config, seed=config.seed, size=config.size) + self.ARC_1D_TASKS = ARC_1D_TASKS self.task_names = list(ARC_1D_TASKS.keys()) def __getitem__(self, idx: int) -> dict: @@ -1179,37 +47,37 @@ def __getitem__(self, idx: int) -> dict: - metadata: dict with generation parameters """ # Create deterministic RNG from base seed and idx - item_rng = random.Random(self.seed + idx) - + item_rng = Random(self.seed + idx) + # Select random task task_name = item_rng.choice(self.task_names) - task_func, task_kwargs = ARC_1D_TASKS[task_name] - + task_func, task_kwargs = self.ARC_1D_TASKS[task_name] + # Generate training examples train_examples = [] size = item_rng.randint(self.config.min_size, self.config.max_size) - + for _ in range(self.config.num_train): example = None while example is None: example = task_func(item_rng, size, **task_kwargs) - + train_examples.append(example) - + # Generate test example test_example = None while test_example is None: test_example = task_func(item_rng, size, **task_kwargs) - + # Format question question = "Find the common rule that maps an input grid to an output grid, given the examples below.\n\n" - + # Add training examples for i, example in enumerate(train_examples, 1): question += f"Example {i}:\n" question += "Input: " + " ".join(str(x) for x in example["input"]) + "\n" question += "Output: " + " ".join(str(x) for x in example["output"]) + "\n\n" - + # Add test input question += "Below is a test input grid. Predict the corresponding output grid by applying the rule you found. " question += "Describe how you derived the rule and your overall reasoning process in detail before you submit your answer. " @@ -1229,33 +97,5 @@ def __getitem__(self, idx: int) -> dict: } - -def task_mirror(task_result: Optional[Dict[str, List[int]]]) -> Optional[Dict[str, List[int]]]: - """Mirror the input and output arrays of a task result.""" - if task_result is None: - return None - return { - "input": list(reversed(task_result["input"])), - "output": list(reversed(task_result["output"])) - } - - -def task_inverse(task_result: Optional[Dict[str, List[int]]]) -> Optional[Dict[str, List[int]]]: - """Swap the input and output arrays of a task result.""" - if task_result is None: - return None - return { - "input": task_result["output"], - "output": task_result["input"] - } - - -def task_identity(task_result: Optional[Dict[str, List[int]]]) -> Optional[Dict[str, List[int]]]: - """Return the task result unchanged.""" - return task_result - - # Register the dataset register_dataset("arc_1d", Arc1DDataset, Arc1DConfig) - - diff --git a/reasoning_gym/cognition/arc_1d_tasks.py b/reasoning_gym/cognition/arc_1d_tasks.py new file mode 100644 index 00000000..9e1334e6 --- /dev/null +++ b/reasoning_gym/cognition/arc_1d_tasks.py @@ -0,0 +1,1145 @@ +from random import Random +from typing import Dict, List, Optional + + +def gen_field(size: int, color: int = 0) -> List[int]: + """Generate a field of given size filled with specified color (default 0).""" + return [color] * size + + +def write_block(pos: int, block: List[int], field: List[int]) -> List[int]: + """Write a block into a field at given position.""" + result = field.copy() + for i, color in enumerate(block): + result[pos + i] = color + return result + + +def task_move_n_pix(rng: Random, size: int, move_pix: int, solid: bool) -> Optional[Dict[str, List[int]]]: + """Generate a task where a block is moved to the right by move_pix pixels.""" + if size <= move_pix + 1: + return None + + block_size = rng.randint(1, size - move_pix - 1) + block_pos = rng.randint(0, size - block_size - move_pix) + + if solid: + color = rng.randint(1, 9) + block = [color] * block_size + else: + block = [rng.randint(1, 9) for _ in range(block_size)] + + question = write_block(block_pos, block, gen_field(size)) + answer = write_block(block_pos + move_pix, block, gen_field(size)) + + return {"input": question, "output": answer} + + +def task_move_n_pix_wrapped(rng: Random, size: int, move_pix: int, solid: bool) -> Optional[Dict[str, List[int]]]: + """Generate a task where a block is moved to the right by move_pix pixels with wrapping.""" + block_size = rng.randint(1, size) + block_pos = rng.randint(0, size) + + if solid: + color = rng.randint(1, 9) + block = [color] * block_size + else: + block = [rng.randint(1, 9) for _ in range(block_size)] + + question = gen_field(size) + answer = gen_field(size) + + for i, color in enumerate(block): + question[(block_pos + i) % size] = color + answer[(block_pos + move_pix + i) % size] = color + + return {"input": question, "output": answer} + + +def task_gravity(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where all non-zero elements are attracted to the left.""" + density = 0.5 + question = [rng.randint(1, 9) if rng.random() < density else 0 for _ in range(size)] + + non_zero = [x for x in question if x != 0] + answer = non_zero + [0] * (size - len(non_zero)) + + return {"input": question, "output": answer} + + +def task_gravity_counting(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where non-zero elements are counted and represented as a sequence of 1s.""" + density = 0.5 + question = [rng.randint(1, 9) if rng.random() < density else 0 for _ in range(size)] + + count = sum(1 for x in question if x != 0) + answer = [1] * count + [0] * (size - count) + + return {"input": question, "output": answer} + + +def task_gravity_antigravity(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where color 1 moves right and color 2 moves left.""" + density = 0.5 + question = [rng.randint(1, 2) if rng.random() < density else 0 for _ in range(size)] + + color1 = [x for x in question if x == 1] + color2 = [x for x in question if x == 2] + answer = [2] * len(color2) + [0] * (size - len(color1) - len(color2)) + [1] * len(color1) + + return {"input": question, "output": answer} + + +def task_block_touch_dot(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where a block moves to touch (but not cover) a dot.""" + dot_color = 1 + block_color = rng.randint(2, 9) + + block_size = rng.randint(1, size) + dot_pos = rng.randint(0, size) + + can_place_left = dot_pos >= block_size + can_place_right = dot_pos + block_size < size + + if not (can_place_left or can_place_right): + return None + + if can_place_left and can_place_right: + side = rng.choice(["left", "right"]) + elif can_place_left: + side = "left" + else: + side = "right" + + if side == "left": + q_block_pos = rng.randint(0, dot_pos - block_size) + a_block_pos = dot_pos - block_size + else: + q_block_pos = rng.randint(dot_pos + 1, size - block_size) + a_block_pos = dot_pos + 1 + + question = gen_field(size) + question[dot_pos] = dot_color + question = write_block(q_block_pos, [block_color] * block_size, question) + + answer = gen_field(size) + answer[dot_pos] = dot_color + answer = write_block(a_block_pos, [block_color] * block_size, answer) + + return {"input": question, "output": answer} + + +def task_block_touch_dot_n_pix(rng: Random, size: int, move_pix: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where a block moves move_pix pixels toward a dot.""" + dot_color = 2 + block_color = rng.randint(3, 9) + + block_size = rng.randint(1, size) + dot_pos = rng.randint(0, size) + + can_place_left = dot_pos >= block_size + can_place_right = dot_pos + block_size < size + + if not (can_place_left or can_place_right): + return None + + if can_place_left and can_place_right: + side = rng.choice(["left", "right"]) + elif can_place_left: + side = "left" + else: + side = "right" + + if side == "left": + q_block_pos = rng.randint(0, dot_pos - block_size) + distance = (dot_pos - block_size) - q_block_pos + move = min(distance, move_pix) + a_block_pos = q_block_pos + move + else: + q_block_pos = rng.randint(dot_pos + 1, size - block_size) + distance = q_block_pos - (dot_pos + 1) + move = min(distance, move_pix) + a_block_pos = q_block_pos - move + + question = gen_field(size) + question[dot_pos] = dot_color + question = write_block(q_block_pos, [block_color] * block_size, question) + + answer = gen_field(size) + answer[dot_pos] = dot_color + answer = write_block(a_block_pos, [block_color] * block_size, answer) + + return {"input": question, "output": answer} + + +def task_block_scale_to_dot(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where a block scales to touch a dot (keeping one end fixed).""" + dot_color = 2 + block_color = rng.randint(3, 9) + + block_size = rng.randint(1, size) + dot_pos = rng.randint(0, size) + + can_place_left = dot_pos >= block_size + can_place_right = dot_pos + block_size < size + + if not (can_place_left or can_place_right): + return None + + if can_place_left and can_place_right: + side = rng.choice(["left", "right"]) + elif can_place_left: + side = "left" + else: + side = "right" + + if side == "left": + q_block_pos = rng.randint(0, dot_pos - block_size) + new_size = dot_pos - q_block_pos + 1 + a_block_pos = q_block_pos + else: + q_block_pos = rng.randint(dot_pos + 1, size - block_size) + new_size = (q_block_pos + block_size) - dot_pos + a_block_pos = dot_pos + + question = gen_field(size) + question[dot_pos] = dot_color + question = write_block(q_block_pos, [block_color] * block_size, question) + + answer = gen_field(size) + answer[dot_pos] = dot_color + answer = write_block(a_block_pos, [block_color] * new_size, answer) + + return {"input": question, "output": answer} + + +def task_two_points_and_fill(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where space between two points of same color is filled with that color.""" + color = rng.randint(1, 9) + + pos1 = rng.randint(0, size - 1) + pos2 = rng.randint(0, size - 1) + if pos1 == pos2: + return None + + pos1, pos2 = min(pos1, pos2), max(pos1, pos2) + + question = gen_field(size) + question[pos1] = color + question[pos2] = color + + answer = question.copy() + for i in range(pos1, pos2 + 1): + answer[i] = color + + return {"input": question, "output": answer} + + +def task_reflect_block_with_border_pixel(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where a block with a border pixel is reflected.""" + block_size = rng.randint(2, size) + if block_size > size: + return None + + c1 = rng.randint(1, 9) + c2 = rng.randint(1, 9) + if c1 == c2: + return None + + side = "left" if rng.random() < 0.5 else "right" + pos = rng.randint(0, size - block_size) + + block = [c1] * block_size + if side == "left": + block[0] = c2 + else: + block[block_size - 1] = c2 + + question = write_block(pos, block, gen_field(size)) + reversed_block = block[::-1] # Reverse the block + answer = write_block(pos, reversed_block, gen_field(size)) + + return {"input": question, "output": answer} + + +def task_reflect_block_with_border_pixel_random(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where a random-colored block with a border pixel is reflected.""" + block_size = rng.randint(2, size) + if block_size > size: + return None + + side = "left" if rng.random() < 0.5 else "right" + pos = rng.randint(0, size - block_size) + + block = [rng.randint(1, 9) for _ in range(block_size)] + border_color = rng.randint(1, 9) + + if side == "left": + if block[0] == border_color: + return None + block[0] = border_color + else: + if block[block_size - 1] == border_color: + return None + block[block_size - 1] = border_color + + question = write_block(pos, block, gen_field(size)) + reversed_block = block[::-1] # Reverse the block + answer = write_block(pos, reversed_block, gen_field(size)) + + return {"input": question, "output": answer} + + +def task_reflect_block_around_dot(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where a block is reflected around a dot.""" + dot_color = 2 + + dot_pos = rng.randint(0, size) + block_size = rng.randint(1, size) + block_pos = rng.randint(0, size - block_size) + block_end = block_pos + block_size - 1 + + # Check if block is strictly to left or right of dot + strictly_left = block_end < dot_pos + strictly_right = block_pos > dot_pos + + if not (strictly_left or strictly_right): + return None + + block_color = rng.randint(3, 9) # Different from dot color + block = [block_color] * block_size + + # Calculate reflection bounds + min_reflect = 2 * dot_pos - block_end + max_reflect = 2 * dot_pos - block_pos + if min_reflect < 0 or max_reflect >= size: + return None + + question = gen_field(size) + question = write_block(block_pos, block, question) + question[dot_pos] = dot_color + + answer = gen_field(size) + answer[dot_pos] = dot_color + for i in range(block_size): + reflect_idx = 2 * dot_pos - (block_pos + i) + answer[reflect_idx] = block[i] + + return {"input": question, "output": answer} + + +def task_block_and_noise_remove(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where noise around a block needs to be removed.""" + block_size = rng.randint(2, size) + if block_size > size: + return None + + block_pos = rng.randint(0, size - block_size) + color = rng.randint(1, 9) + + # Create field with block + field = gen_field(size) + for i in range(block_size): + field[block_pos + i] = color + + # Track forbidden positions for noise + forbidden = [False] * size + for i in range(block_pos, block_pos + block_size): + forbidden[i] = True + if block_pos > 0: + forbidden[block_pos - 1] = True + if block_pos + block_size < size: + forbidden[block_pos + block_size] = True + + # Add noise + noise_count = rng.randint(1, 3) + noise_positions = [] + + for _ in range(noise_count): + allowed = [i for i in range(size) if not forbidden[i]] + if not allowed: + break + noise_pos = rng.choice(allowed) + noise_positions.append(noise_pos) + field[noise_pos] = color + forbidden[noise_pos] = True + if noise_pos > 0: + forbidden[noise_pos - 1] = True + if noise_pos + 1 < size: + forbidden[noise_pos + 1] = True + + if len(noise_positions) < noise_count: + return None + + question = field + answer = field.copy() + for pos in noise_positions: + answer[pos] = 0 + + return {"input": question, "output": answer} + + +def task_block_and_noise_remove_inside(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where noise inside a block needs to be removed.""" + if size <= 6: + return None + + block_size = rng.randint(6, size) + if block_size > size: + return None + + block_pos = rng.randint(0, size - block_size) + color = rng.randint(1, 9) + + # Create field with block + field = gen_field(size) + for i in range(block_size): + field[block_pos + i] = color + + # Add noise inside block + max_noise = max(1, (block_size // 2) - 1) + noise_count = rng.randint(1, max_noise) + + positions = list(range(block_size)) + rng.shuffle(positions) + noise_positions = positions[:noise_count] + + for offset in noise_positions: + pos = block_pos + offset + noise_color = rng.randint(1, 9) + while noise_color == color: + noise_color = rng.randint(1, 9) + field[pos] = noise_color + + question = field + answer = field.copy() + for offset in noise_positions: + answer[block_pos + offset] = color + + return {"input": question, "output": answer} + + +def task_copy_block_to_dots(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where a block pattern is copied to dot positions.""" + block_size = 3 if rng.random() < 0.5 else 5 + if block_size >= size: + return None + + color = rng.randint(1, 9) + block = [color] * block_size + + # Generate dots with minimum distance to prevent overlap + min_gap = block_size + dot_positions = [] + pos = block_size + block_size // 2 + 1 + + while pos <= size - block_size: + if rng.random() < 0.5: # Control dot density + dot_positions.append(pos) + pos += min_gap + pos += 1 + + if not dot_positions: + return None + + question = gen_field(size) + question = write_block(0, block, question) + for pos in dot_positions: + question[pos] = color + + answer = gen_field(size) + answer = write_block(0, block, answer) + for pos in dot_positions: + block_start = pos - block_size // 2 + answer = write_block(block_start, block, answer) + + return {"input": question, "output": answer} + + +def task_copy_block_to_dots_colors(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where a block pattern is copied to dot positions with matching colors.""" + block_size = 3 if rng.random() < 0.5 else 5 + if block_size >= size: + return None + + block_color = rng.randint(1, 9) + block = [block_color] * block_size + + # Generate dots with minimum distance to prevent overlap + min_gap = block_size + dot_positions = [] + dot_colors = [] + pos = block_size + block_size // 2 + 1 + + while pos < size - block_size: + if rng.random() < 0.5: + dot_color = rng.randint(1, 9) + dot_positions.append(pos) + dot_colors.append(dot_color) + pos += min_gap + pos += 1 + + if not dot_positions: + return None + + question = gen_field(size) + question = write_block(0, block, question) + for i, pos in enumerate(dot_positions): + question[pos] = dot_colors[i] + + answer = gen_field(size) + answer = write_block(0, block, answer) + for i, pos in enumerate(dot_positions): + block_start = pos - block_size // 2 + colored_block = [dot_colors[i]] * block_size + answer = write_block(block_start, colored_block, answer) + + return {"input": question, "output": answer} + + +def task_paint_biggest_block(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where the largest block is painted a different color.""" + target_color = 1 + initial_color = rng.randint(2, 9) + + # Generate random blocks + question = gen_field(size) + blocks = [] + pos = 0 + + while pos < size: + if rng.random() < 0.4 and size - pos >= 2: + block_size = rng.randint(2, min(size - pos, 6)) + blocks.append((pos, block_size)) + for i in range(block_size): + question[pos + i] = initial_color + pos += block_size + 1 + else: + pos += 1 + + if len(blocks) < 2: + return None + + # Find biggest block + biggest_pos, biggest_size = max(blocks, key=lambda x: x[1]) + + # Check if there are multiple blocks of the same size + biggest_count = sum(1 for _, size in blocks if size == biggest_size) + if biggest_count > 1: + return None + + answer = question.copy() + for i in range(biggest_size): + answer[biggest_pos + i] = target_color + + return {"input": question, "output": answer} + + +def task_sort_blocks_by_size(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where blocks are sorted by size with 1 pixel gaps.""" + color = rng.randint(1, 9) + blocks = [] + pos = 0 + + # Generate random blocks with random sizes + while pos < size: + if rng.random() < 0.4 and size - pos >= 2: + block_size = rng.randint(1, min(size - pos, 6)) + blocks.append((pos, block_size)) + pos += block_size + rng.randint(1, 4) # Random gaps + else: + pos += 1 + + if len(blocks) < 2: + return None + + # Create input field + question = gen_field(size) + for pos, block_size in blocks: + for i in range(block_size): + question[pos + i] = color + + # Sort blocks by size + blocks.sort(key=lambda x: x[1]) + + # Check if sorted blocks fit with gaps + total_space = sum(size for _, size in blocks) + len(blocks) - 1 + if total_space > size: + return None + + # Create answer field with sorted blocks + answer = gen_field(size) + current_pos = 0 + + for _, block_size in blocks: + for i in range(block_size): + answer[current_pos + i] = color + current_pos += block_size + 1 # One pixel gap + + return {"input": question, "output": answer} + + +def task_sort_complete_sequence(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where a complete sequence of block sizes is sorted.""" + # Calculate max possible block size given total array size + max_size = 1 + total_space = 0 + while total_space + max_size + 1 <= size: + total_space += max_size + 1 + max_size += 1 + max_size -= 1 + + if max_size < 2: + return None + + color = rng.randint(1, 9) + + # Create sequence of all sizes from 1 to max_size + blocks = list(range(1, max_size + 1)) + rng.shuffle(blocks) + + # Create input field with shuffled blocks + question = gen_field(size) + pos = 0 + for block_size in blocks: + for i in range(block_size): + question[pos + i] = color + pos += block_size + 1 + + # Create answer field with sorted blocks + answer = gen_field(size) + pos = 0 + for block_size in range(1, max_size + 1): + for i in range(block_size): + answer[pos + i] = color + pos += block_size + 1 + + return {"input": question, "output": answer} + + +def task_recolor_blocks_by_size(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where two blocks are recolored based on their size.""" + # Generate two different random sizes + size1 = rng.randint(2, 8) + size2 = rng.randint(2, 8) + while size2 == size1: + size2 = rng.randint(2, 8) + + # Ensure both blocks fit with at least 1 gap + if size1 + size2 + 1 > size: + return None + + # Place blocks with gap + pos1 = rng.randint(0, size - (size1 + size2 + 1)) + pos2 = rng.randint(pos1 + size1 + 1, size - size2) + + # Create input field with both blocks color 3 + question = gen_field(size) + for i in range(size1): + question[pos1 + i] = 3 + for i in range(size2): + question[pos2 + i] = 3 + + # Create answer field with recolored blocks + answer = question.copy() + if size1 > size2: + for i in range(size1): + answer[pos1 + i] = 1 + for i in range(size2): + answer[pos2 + i] = 2 + else: + for i in range(size1): + answer[pos1 + i] = 2 + for i in range(size2): + answer[pos2 + i] = 1 + + return {"input": question, "output": answer} + + +def task_gravity_one_step(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where non-zero elements move one step left if possible.""" + question = [rng.randint(1, 9) if rng.random() < 0.5 else 0 for _ in range(size)] + answer = question.copy() + + # Move each non-zero pixel one step left if possible + for i in range(1, size): + if answer[i] != 0 and answer[i - 1] == 0: + answer[i - 1] = answer[i] + answer[i] = 0 + + return {"input": question, "output": answer} + + +def task_move_block_by_own_size(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where a block moves right by its own size.""" + block_size = rng.randint(1, size // 2) # Ensure space for movement + pos = rng.randint(0, size - block_size * 2) # Space for block and movement + color = rng.randint(1, 9) + + question = gen_field(size) + block = [color] * block_size + question = write_block(pos, block, question) + + answer = write_block(pos + block_size, block, gen_field(size)) + + return {"input": question, "output": answer} + + +def task_change_to_five(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where all non-zero colors change to 5.""" + density = 0.5 + question = [rng.randint(1, 9) if rng.random() < density else 0 for _ in range(size)] + answer = [5 if x != 0 else 0 for x in question] + + return {"input": question, "output": answer} + + +def task_recolor_blocks_from_palette(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where blocks are recolored using a color palette.""" + # Generate blocks of same size + block_size = rng.randint(2, 4) + blocks = [] + pos = 0 + + while pos + block_size <= size: + if rng.random() < 0.4: + blocks.append(pos) + pos += block_size + 1 + else: + pos += 1 + + # Ensure we have space for palette + while blocks and blocks[-1] + block_size + len(blocks) + 1 >= size: + blocks.pop() + + if not blocks: + return None + + # Shift blocks right to make room for palette + palette_size = len(blocks) + blocks = [pos + palette_size + 1 for pos in blocks] + + # Generate color palette + colors = [] + for _ in range(len(blocks)): + while True: + color = rng.randint(1, 9) + if color not in colors: + colors.append(color) + break + + # Create question with color palette and blocks + question = gen_field(size) + + # Place color palette at start + for i, color in enumerate(colors): + question[i] = color + + # Place blocks of color 5 + for block_pos in blocks: + for i in range(block_size): + question[block_pos + i] = 5 + + # Create answer with recolored blocks + answer = question.copy() + for block_idx, block_pos in enumerate(blocks): + color = colors[block_idx] + for i in range(block_size): + answer[block_pos + i] = color + + return {"input": question, "output": answer} + + +def task_duplicate_block_from_seeds(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where a block is duplicated from seed pixels.""" + block_size = rng.randint(2, 4) + if block_size + 1 >= size: + return None + if size <= 3 + block_size: + return None + + # Position block with space for seeds + block_pos = rng.randint(2, size - block_size - 1) + + # Decide seed placement + left_seed = rng.random() < 0.5 + right_seed = rng.random() < 0.5 + if not (left_seed or right_seed): + return None + + # Create input + question = gen_field(size) + + # Place main block + for i in range(block_size): + question[block_pos + i] = 1 + + # Place seeds with gaps + seeds = [] + if left_seed: + color = rng.randint(1, 9) + question[block_pos - 2] = color + seeds.append(("left", block_pos - 2, color)) + if right_seed: + color = rng.randint(1, 9) + question[block_pos + block_size + 1] = color + seeds.append(("right", block_pos + block_size + 1, color)) + + # Create answer with duplicated blocks + answer = question.copy() + + for side, seed_pos, color in seeds: + if side == "left": + # For left seed, blocks end at seed + end_pos = seed_pos + while end_pos >= 0: + start_pos = end_pos - block_size + 1 + for pos in range(max(0, start_pos), end_pos + 1): + answer[pos] = color + if start_pos < 1: + break + end_pos = start_pos - 2 # -1 for gap + else: # side == "right" + # For right seed, blocks start at seed + start_pos = seed_pos + while start_pos < size: + for offset in range(min(block_size, size - start_pos)): + answer[start_pos + offset] = color + if start_pos + block_size + 1 >= size: + break + start_pos = start_pos + block_size + 1 # +1 for gap + + return {"input": question, "output": answer} + + +def task_fill_from_pixel(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where a pixel fills in one direction until hitting another pixel.""" + block_size = rng.randint(3, 6) + if block_size >= size - 2: + return None + + # Position block with space for seed + block_pos = rng.randint(1, size - block_size - 1) + + # Create input + question = gen_field(size) + + # Place main block + block_color = rng.randint(1, 9) + for i in range(block_size): + question[block_pos + i] = block_color + + # Place seed pixel and determine fill direction + seed_color = rng.randint(1, 9) + while seed_color == block_color: + seed_color = rng.randint(1, 9) + + is_left = rng.random() < 0.5 + + if is_left: + question[block_pos - 1] = seed_color + else: + question[block_pos + block_size] = seed_color + + # Create answer with fill + answer = question.copy() + + if is_left: + # Fill from seed to left border + for i in range(block_pos): + answer[i] = seed_color + else: + # Fill from seed to right border + for i in range(block_pos + block_size, size): + answer[i] = seed_color + + return {"input": question, "output": answer} + + +def task_mark_size_two_blocks(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where size-2 blocks are marked with surrounding pixels.""" + blocks = [] + pos = 0 + + # Generate blocks with minimum gap of 2 + while pos < size: + if rng.random() < 0.4: + block_size = rng.randint(1, 3) + # Check if we have space for block and potential markers + needed_space = block_size + (2 if block_size == 2 else 0) + if pos + needed_space < size: + blocks.append((pos, block_size)) + pos += block_size + 2 # Minimum gap of 2 + + pos += 1 + + if len(blocks) < 2: + return None + + # Verify gaps between blocks (including markers) + valid = True + for i in range(len(blocks) - 1): + pos1, size1 = blocks[i] + pos2, _ = blocks[i + 1] + needed_gap = 3 if size1 == 2 else 2 + if pos2 - (pos1 + size1) < needed_gap: + valid = False + break + if not valid: + return None + + # Create input with blocks + question = gen_field(size) + for pos, block_size in blocks: + # Place block + for i in range(block_size): + question[pos + i] = 1 + + # Create answer with markers + answer = question.copy() + for pos, block_size in blocks: + if block_size == 2: + # Add markers for size 2 blocks + if pos > 0: + answer[pos - 1] = 3 + if pos + block_size < size: + answer[pos + block_size] = 3 + + return {"input": question, "output": answer} + + +def task_fill_until_collision(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where pixels fill empty space until collision.""" + # At least 4 positions for meaningful puzzle + if size < 4: + return None + + is_left = rng.random() < 0.5 + question = gen_field(size) + + # Place the side marker + if is_left: + question[0] = 5 + else: + question[size - 1] = 5 + + # Place 2-4 random pixels + num_pixels = rng.randint(2, 4) + positions = [] + + if is_left: + # Skip first position + for _ in range(num_pixels): + while True: + pos = rng.randint(1, size - 1) + if pos not in positions: + positions.append(pos) + break + else: + # Skip last position + for _ in range(num_pixels): + while True: + pos = rng.randint(0, size - 2) + if pos not in positions: + positions.append(pos) + break + + # Color random pixels + for pos in positions: + question[pos] = rng.randint(1, 9) + + positions.sort() + + # Create answer + answer = question.copy() + + if is_left: + # Fill right from each pixel + prev_pos = 0 # Start from marker + for pos in positions: + color = question[pos] + # Fill from previous position to current + for i in range(prev_pos + 1, pos): + answer[i] = color + prev_pos = pos + else: + # Fill left from each pixel + prev_pos = size - 1 # Start from marker + for pos in reversed(positions): + color = question[pos] + # Fill from current position to previous + for i in range(pos + 1, prev_pos): + answer[i] = color + prev_pos = pos + + return {"input": question, "output": answer} + + +def task_repeat_pattern_full(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where a pattern is repeated to fill the space.""" + # Generate initial pattern + pattern_size = rng.randint(2, 5) + pattern = [rng.randint(1, 9) for _ in range(pattern_size)] + + # Calculate total size needed for 2 repetitions + double_size = pattern_size * 2 + if double_size >= size: + return None + + # Create input with 2 repetitions + question = gen_field(size) + for i in range(pattern_size): + question[i] = pattern[i] + question[i + pattern_size] = pattern[i] + + # Create answer with maximum repetitions + answer = gen_field(size) + pos = 0 + while pos + pattern_size <= size: + for i in range(pattern_size): + answer[pos + i] = pattern[i] + pos += pattern_size + + # Fill remaining space (if any) with pattern elements + for i in range(pos, size): + answer[i] = pattern[i - pos] + + return {"input": question, "output": answer} + + +def task_gravity_weighted_colors(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where color 2 is heavier than color 1 in gravity.""" + # Generate random field with only colors 1 and 2 + question = [rng.randint(1, 2) if rng.random() < 0.5 else 0 for _ in range(size)] + + # Count colors + count_1 = sum(1 for x in question if x == 1) + count_2 = sum(1 for x in question if x == 2) + + # Create answer with sorted colors + answer = gen_field(size) + + # Place heavier color 2 first + for i in range(count_2): + answer[i] = 2 + + # Then place color 1 + for i in range(count_1): + answer[count_2 + i] = 1 + + return {"input": question, "output": answer} + + +def task_color_left_half_blocks(rng: Random, size: int) -> Optional[Dict[str, List[int]]]: + """Generate a task where left half of blocks are colored differently.""" + pos = 0 + question = gen_field(size) + blocks = [] + + # Generate blocks with gap 1 + while pos < size: + if rng.random() < 0.4: + block_size = rng.randint(2, 8) + if pos + block_size >= size: + break + + blocks.append((pos, block_size)) + for i in range(block_size): + question[pos + i] = 2 + pos += block_size + 1 # block size + gap + else: + pos += 1 + + if len(blocks) < 2: + return None + + # Create answer with half-colored blocks + answer = question.copy() + for pos, block_size in blocks: + half_size = block_size // 2 + for i in range(half_size): + answer[pos + i] = 8 + + return {"input": question, "output": answer} + + +def task_mirror(task_result: Optional[Dict[str, List[int]]]) -> Optional[Dict[str, List[int]]]: + """Mirror the input and output arrays of a task result.""" + if task_result is None: + return None + return {"input": list(reversed(task_result["input"])), "output": list(reversed(task_result["output"]))} + + +def task_inverse(task_result: Optional[Dict[str, List[int]]]) -> Optional[Dict[str, List[int]]]: + """Swap the input and output arrays of a task result.""" + if task_result is None: + return None + return {"input": task_result["output"], "output": task_result["input"]} + + +def task_identity(task_result: Optional[Dict[str, List[int]]]) -> Optional[Dict[str, List[int]]]: + """Return the task result unchanged.""" + return task_result + + +# Table of all ARC 1D task functions with their parameters +ARC_1D_TASKS = { + # Move tasks + "move_1pix_solid": (task_move_n_pix, {"move_pix": 1, "solid": True}), + "move_2pix_solid": (task_move_n_pix, {"move_pix": 2, "solid": True}), + "move_3pix_solid": (task_move_n_pix, {"move_pix": 3, "solid": True}), + "move_4pix_solid": (task_move_n_pix, {"move_pix": 4, "solid": True}), + "move_1pix_colorful": (task_move_n_pix, {"move_pix": 1, "solid": False}), + "move_2pix_colorful": (task_move_n_pix, {"move_pix": 2, "solid": False}), + "move_3pix_colorful": (task_move_n_pix, {"move_pix": 3, "solid": False}), + "move_4pix_colorful": (task_move_n_pix, {"move_pix": 4, "solid": False}), + # Move wrapped tasks + "move_1pix_solid_wrapped": (task_move_n_pix_wrapped, {"move_pix": 1, "solid": True}), + "move_2pix_solid_wrapped": (task_move_n_pix_wrapped, {"move_pix": 2, "solid": True}), + "move_3pix_solid_wrapped": (task_move_n_pix_wrapped, {"move_pix": 3, "solid": True}), + "move_4pix_solid_wrapped": (task_move_n_pix_wrapped, {"move_pix": 4, "solid": True}), + "move_1pix_colorful_wrapped": (task_move_n_pix_wrapped, {"move_pix": 1, "solid": False}), + "move_2pix_colorful_wrapped": (task_move_n_pix_wrapped, {"move_pix": 2, "solid": False}), + "move_3pix_colorful_wrapped": (task_move_n_pix_wrapped, {"move_pix": 3, "solid": False}), + "move_4pix_colorful_wrapped": (task_move_n_pix_wrapped, {"move_pix": 4, "solid": False}), + # Gravity tasks + "gravity": (task_gravity, {}), + "gravity_counting": (task_gravity_counting, {}), + "gravity_antigravity": (task_gravity_antigravity, {}), + "gravity_one_step": (task_gravity_one_step, {}), + "gravity_weighted_colors": (task_gravity_weighted_colors, {}), + # Block tasks + "block_touch_dot": (task_block_touch_dot, {}), + "block_touch_dot_1pix": (task_block_touch_dot_n_pix, {"move_pix": 1}), + "block_touch_dot_2pix": (task_block_touch_dot_n_pix, {"move_pix": 2}), + "block_touch_dot_3pix": (task_block_touch_dot_n_pix, {"move_pix": 3}), + "block_touch_dot_4pix": (task_block_touch_dot_n_pix, {"move_pix": 4}), + "block_scale_to_dot": (task_block_scale_to_dot, {}), + "block_and_noise_remove": (task_block_and_noise_remove, {}), + "block_and_noise_remove_inside": (task_block_and_noise_remove_inside, {}), + "move_block_by_own_size": (task_move_block_by_own_size, {}), + # Pattern tasks + "two_points_and_fill": (task_two_points_and_fill, {}), + "copy_block_to_dots": (task_copy_block_to_dots, {}), + "copy_block_to_dots_colors": (task_copy_block_to_dots_colors, {}), + "repeat_pattern_full": (task_repeat_pattern_full, {}), + # Reflection tasks + "reflect_block_with_border_pixel": (task_reflect_block_with_border_pixel, {}), + "reflect_block_random": (task_reflect_block_with_border_pixel_random, {}), + "reflect_block_around_dot": (task_reflect_block_around_dot, {}), + # Color tasks + "paint_biggest_block": (task_paint_biggest_block, {}), + "recolor_blocks_by_size": (task_recolor_blocks_by_size, {}), + "change_to_five": (task_change_to_five, {}), + "recolor_blocks_from_palette": (task_recolor_blocks_from_palette, {}), + "color_left_half_blocks": (task_color_left_half_blocks, {}), + # Sorting tasks + "sort_blocks_by_size": (task_sort_blocks_by_size, {}), + "sort_complete_sequence": (task_sort_complete_sequence, {}), + # Fill tasks + "duplicate_block_from_seeds": (task_duplicate_block_from_seeds, {}), + "fill_from_pixel": (task_fill_from_pixel, {}), + "fill_until_collision": (task_fill_until_collision, {}), + # Marking tasks + "mark_size_two_blocks": (task_mark_size_two_blocks, {}), +} diff --git a/tests/test_arithmetic.py b/tests/test_basic_arithmetic.py similarity index 100% rename from tests/test_arithmetic.py rename to tests/test_basic_arithmetic.py