diff --git a/reasoning_gym/algorithmic/__init__.py b/reasoning_gym/algorithmic/__init__.py index fe0a2dc2..875ab539 100644 --- a/reasoning_gym/algorithmic/__init__.py +++ b/reasoning_gym/algorithmic/__init__.py @@ -19,6 +19,7 @@ from .manipulate_matrix import ManipulateMatrixConfig, ManipulateMatrixDataset from .number_filtering import NumberFilteringConfig, NumberFilteringDataset from .number_sorting import NumberSortingConfig, NumberSortingDataset from .palindrome_generation import PalindromeConfig, PalindromeDataset +from .pool_matrix import PoolMatrixConfig, PoolMatrixDataset from .ransom_note import RansomNoteConfig, RansomNoteDataset from .rotate_matrix import RotateMatrixConfig, RotateMatrixDataset from .sentence_reordering import SentenceReorderingConfig, SentenceReorderingDataset @@ -68,6 +69,8 @@ __all__ = [ "ManipulateMatrixDataset", "BinaryMatrixConfig", "BinaryMatrixDataset", + "PoolMatrixConfig", + "PoolMatrixDataset", "ABConfig", "ABDataset", "CountPrimesConfig", diff --git a/reasoning_gym/algorithmic/pool_matrix.py b/reasoning_gym/algorithmic/pool_matrix.py new file mode 100644 index 00000000..dda7ed2d --- /dev/null +++ b/reasoning_gym/algorithmic/pool_matrix.py @@ -0,0 +1,142 @@ +"""Perform average / max pooling on a matrix""" + +from copy import deepcopy +from dataclasses import dataclass +from random import Random +from typing import Dict, Optional + +import numpy as np + +from ..factory import ProceduralDataset, register_dataset + +QUESTION_TEMPLATE = """Your job is to perform max/average pooling on the given matrix. +The stride is equal to the kernel size, meaning there is no overlap between the pooling regions. + +Example 1: +- Input: Perform max pooling on the following matrix with a kernel size of 2: +1 2 3 4 +5 6 7 8 +9 10 11 12 +13 14 15 16 +- Output: +6 8 +14 16 + +Example 2: +- Input: Perform average pooling on the following matrix with a kernel size of 2: +1 2 3 4 +5 6 7 8 +9 10 11 12 +13 14 15 16 +- Output: +3.5 5.5 +11.5 13.5 + +Perform {pool_type} pooling on the following matrix with a kernel size of {pool_size}: +{matrix} +""" + + +@dataclass +class PoolMatrixConfig: + """Configuration for Pool Matrix dataset generation""" + + min_rows: int = 2 # Minimum rows of the matrix + min_cols: int = 2 # Minimum columns of the matrix + max_rows: int = 10 # Maximum rows of the matrix + max_cols: int = 10 # Maximum columns of the matrix + max_pool_size: int = 3 # Maximum pooling size + + size: int = 500 # Virtual dataset size + seed: Optional[int] = None + + def validate(self): + """Validate configuration parameters""" + assert 2 <= self.min_rows, "min_rows must be at least 2" + assert 2 <= self.min_cols, "min_cols must be at least 2" + assert self.min_rows <= self.max_rows, "max_rows must be at least min_rows" + assert self.min_cols <= self.max_cols, "max_cols must be at least min_cols" + assert 1 <= self.max_pool_size, "max_pool_size must be at least 1" + + +class PoolMatrixDataset(ProceduralDataset): + """Generates Pool Matrix exercises with configurable difficulty""" + + def __init__(self, config: PoolMatrixConfig): + super().__init__(config=config, seed=config.seed, size=config.size) + + def _get_matrix(self, rng: Random) -> np.ndarray: + """Generate a random matrix""" + rows = rng.randint(self.config.min_rows, self.config.max_rows) + cols = rng.randint(self.config.min_rows, self.config.max_cols) + return np.random.randint(0, 10, (rows, cols)) + + def _matrix_to_str(self, matrix: np.ndarray) -> str: + """Get a string representation of the matrix""" + return "\n".join(" ".join(str(round(x, 2)) for x in row) for row in matrix) + + def _max_pool(self, matrix: np.ndarray, pool_size: int) -> np.ndarray: + """Perform max pooling on the matrix""" + rows, cols = matrix.shape + return np.array( + [ + [np.max(matrix[i : i + pool_size, j : j + pool_size]) for j in range(0, cols, pool_size)] + for i in range(0, rows, pool_size) + ] + ) + + def _average_pool(self, matrix: np.ndarray, pool_size: int) -> np.ndarray: + """Perform average pooling on the matrix""" + rows, cols = matrix.shape + return np.array( + [ + [np.mean(matrix[i : i + pool_size, j : j + pool_size]) for j in range(0, cols, pool_size)] + for i in range(0, rows, pool_size) + ] + ) + + def score_answer(self, answer: Optional[str], entry: Dict[str, any]) -> float: + """Score the answer based on the metadata""" + + reward = 0.0 + try: + if answer is not None: + oracle_answer = np.array(entry["answer"]) + answer = np.array(answer) + if oracle_answer.shape == answer.shape and np.allclose(oracle_answer, answer): + reward = 1.0 + if oracle_answer.shape == answer.shape: + reward = 0.1 + else: + reward = 0.01 + except: + print("Error in scoring answer for Pool Matrix") + return reward + + def __getitem__(self, idx: int) -> dict: + """Generate a single Pool Matrix question""" + rng = Random(self.seed + idx) + np.random.seed(self.seed + idx) + + matrix = self._get_matrix(rng) + matrix_str = self._matrix_to_str(matrix) + + pool_size = rng.randint(1, self.config.max_pool_size) + pool_type = rng.choice(["average", "max"]) + + answer = self._average_pool(matrix, pool_size) if pool_type == "average" else self._max_pool(matrix, pool_size) + answer_str = self._matrix_to_str(answer) + + return { + "question": QUESTION_TEMPLATE.format(matrix=matrix_str, pool_type=pool_type, pool_size=pool_size), + "answer": answer_str, + "metadata": { + "matrix": matrix.tolist(), + "pool_type": pool_type, + "pool_size": pool_size, + "solution": answer.tolist(), + }, + } + + +register_dataset("pool_matrix", PoolMatrixDataset, PoolMatrixConfig) diff --git a/tests/test_pool_matrix.py b/tests/test_pool_matrix.py new file mode 100644 index 00000000..aa3fe6b6 --- /dev/null +++ b/tests/test_pool_matrix.py @@ -0,0 +1,138 @@ +"""Tests for Pool Matrix questions generation""" + +import numpy as np +import pytest + +from reasoning_gym.algorithmic.pool_matrix import PoolMatrixConfig, PoolMatrixDataset + + +def test_pool_matrix_config_validation(): + """Test that invalid configs raise appropriate errors""" + + for field in ["min_rows", "min_cols", "max_rows", "max_cols"]: + with pytest.raises(AssertionError): + config = PoolMatrixConfig(**{field: -1}) # Negative not allowed + config.validate() + + with pytest.raises(AssertionError): + config = PoolMatrixConfig(**{field: 0}) # Zero not allowed + config.validate() + + with pytest.raises(AssertionError): + config = PoolMatrixConfig(**{field: 1}) # One not allowed + config.validate() + + with pytest.raises(AssertionError): + config = PoolMatrixConfig(max_pool_size=-1) # Negative not allowed + config.validate() + + with pytest.raises(AssertionError): + config = PoolMatrixConfig(max_pool_size=0) # Zero not allowed + config.validate() + + +def test_pool_matrix_dataset_deterministic(): + """Test that dataset generates same items with same seed""" + config = PoolMatrixConfig(seed=42, size=10) + dataset1 = PoolMatrixDataset(config) + dataset2 = PoolMatrixDataset(config) + + for i in range(len(dataset1)): + assert dataset1[i] == dataset2[i] + + +def test_pool_matrix_dataset_items(): + """Test basic properties of generated items""" + config = PoolMatrixConfig(max_rows=10, max_cols=10, max_pool_size=3, size=10, seed=42) + dataset = PoolMatrixDataset(config) + + for i in range(len(dataset)): + item = dataset[i] + # Check item structure + assert isinstance(item, dict) + assert "question" in item + assert "answer" in item + assert "metadata" in item + + # Check metadata + assert "matrix" in item["metadata"] + assert "pool_type" in item["metadata"] + assert "pool_size" in item["metadata"] + assert "solution" in item["metadata"] + + matrix = item["metadata"]["matrix"] + pool_type = item["metadata"]["pool_type"] + pool_size = item["metadata"]["pool_size"] + solution = item["metadata"]["solution"] + + # Verify dimensions + assert len(matrix) <= config.max_rows + assert all(len(row) <= config.max_cols for row in matrix) + assert len(solution) <= len(matrix) + assert len(solution[0]) <= len(matrix[0]) + assert pool_size <= config.max_pool_size + assert pool_type in ["average", "max"] + + +def test_pool_matrix_dataset_iteration(): + """Test that iteration respects dataset size""" + config = PoolMatrixConfig(size=5, seed=42) + dataset = PoolMatrixDataset(config) + + items = list(dataset) + assert len(items) == config.size + + # Test multiple iterations yield same items + assert items == list(dataset) + + +def test_pool_matrix_answer(): + """Test the pooling methods""" + config = PoolMatrixConfig(seed=42) + dataset = PoolMatrixDataset(config) + + # 1. Max pooling + matrix = np.array([[1, 2], [3, 4]]) + assert np.allclose(dataset._max_pool(matrix, 2), np.array([[4]])) + + matrix = np.array( + [ + [1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12], + ] + ) + assert np.allclose(dataset._max_pool(matrix, 2), np.array([[6, 8], [10, 12]])) + + matrix = np.array( + [ + [1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12], + [13, 14, 15, 16], + ] + ) + assert np.allclose(dataset._max_pool(matrix, 2), np.array([[6, 8], [14, 16]])) + + # 2. Average pooling + matrix = np.array([[1, 2], [3, 4]]) + assert np.allclose(dataset._average_pool(matrix, 2), np.array([[2.5]])) + + matrix = np.array( + [ + [1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12], + ] + ) + assert np.allclose(dataset._average_pool(matrix, 2), np.array([[3.5, 5.5], [9.5, 11.5]])) + + matrix = np.array( + [ + [1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12], + [13, 14, 15, 16], + ] + ) + assert np.allclose(dataset._average_pool(matrix, 2), np.array([[3.5, 5.5], [11.5, 13.5]]))