Power Function (#102)

* power function dataset + tests

reasoning_gym/arithmetic/__init__.py

@@ -10,6 +10,7 @@ from .gcd import GCDConfig, GCDDataset
 from .gsm_symbolic.gsm_symbolic import GSMSymbolicDataset, GSMSymbolicDatasetConfig
 from .lcm import LCMConfig, LCMDataset
 from .leg_counting import LegCountingConfig, LegCountingDataset
+from .power_function import PowerFunctionConfig, PowerFunctionDataset
 from .prime_factorization import PrimeFactorizationConfig, PrimeFactorizationDataset
 from .time_intervals import TimeIntervalsConfig, TimeIntervalsDataset
 

reasoning_gym/arithmetic/power_function.py

@@ -0,0 +1,62 @@
+"""Computhe the power of a number."""
+
+from dataclasses import dataclass
+from math import pow
+from random import Random
+from typing import Dict, Optional
+
+from ..factory import ProceduralDataset, register_dataset
+
+QUESTION_TEMPLATE = """Compute {base}^{exponent}"""
+
+
+@dataclass
+class PowerFunctionConfig:
+    """Configuration for Power Function dataset generation"""
+
+    min_base: float = -(10**6)  # Minimum base value
+    max_base: float = 10**6  # Maximum base value
+    min_exponent: int = -50  # Minimum exponent value
+    max_exponent: int = 50  # Maximum exponent value
+
+    size: int = 500  # Virtual dataset size
+    seed: Optional[int] = None
+
+
+class PowerFunctionDataset(ProceduralDataset):
+    """Generates Power Function exercises with configurable difficulty"""
+
+    def __init__(self, config: PowerFunctionConfig):
+        super().__init__(config=config, seed=config.seed, size=config.size)
+
+    def score_answer(self, answer: Optional[str], entry: Dict[str, any]) -> float:
+        """Overwrite this method in derived classes if a single oracle answer is not available."""
+        oracle_answer = entry["answer"]
+        reward = 0.0
+        if answer is not None:
+            difference = abs(float(answer) - float(oracle_answer))
+            if difference < 1e-6:
+                reward = 1.0
+            elif difference < 1e-1:
+                reward = 0.5
+            else:
+                reward = 0.01
+
+        return reward
+
+    def __getitem__(self, idx: int) -> dict:
+        """Generate a single Power Function question"""
+        rng = Random(self.seed + idx)
+
+        base = rng.uniform(self.config.min_base, self.config.max_base)
+        exponent = rng.randint(self.config.min_exponent, self.config.max_exponent)
+        answer = pow(base, exponent)
+
+        return {
+            "question": f"Compute {base}^{exponent}",
+            "answer": str(answer),
+            "metadata": {"base": base, "exponent": exponent, "solution": answer},
+        }
+
+
+register_dataset("power_function", PowerFunctionDataset, PowerFunctionConfig)

tests/test_power_function.py

@@ -0,0 +1,78 @@
+"""Tests for Power Function questions generation"""
+
+import pytest
+
+from reasoning_gym.arithmetic import PowerFunctionConfig, PowerFunctionDataset
+
+
+def test_power_function_dataset_deterministic():
+    """Test that dataset generates same items with same seed"""
+    config = PowerFunctionConfig(seed=42, size=10)
+    dataset1 = PowerFunctionDataset(config)
+    dataset2 = PowerFunctionDataset(config)
+
+    for i in range(len(dataset1)):
+        assert dataset1[i] == dataset2[i]
+
+
+def test_power_function_dataset_items():
+    """Test basic properties of generated items"""
+    config = PowerFunctionConfig(min_base=-100, max_base=-100, min_exponent=-10, max_exponent=10, size=10, seed=42)
+    dataset = PowerFunctionDataset(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 "base" in item["metadata"]
+        assert "exponent" in item["metadata"]
+
+        base = item["metadata"]["base"]
+        exponent = item["metadata"]["exponent"]
+        solution = item["metadata"]["solution"]
+
+        # Verify values
+        assert config.min_base <= base <= config.max_base
+        assert config.min_exponent <= exponent <= config.max_exponent
+        assert solution == pow(base, exponent)
+
+
+def test_power_function_dataset_iteration():
+    """Test that iteration respects dataset size"""
+    config = PowerFunctionConfig(size=5, seed=42)
+    dataset = PowerFunctionDataset(config)
+
+    items = list(dataset)
+    assert len(items) == config.size
+
+    # Test multiple iterations yield same items
+    assert items == list(dataset)
+
+
+def test_power_function_score_function():
+    """Test score function"""
+    config = PowerFunctionConfig(seed=42)
+    dataset = PowerFunctionDataset(config)
+
+    item = dataset[0]
+
+    # Answer is within 1e-6 of solution
+    answer = str(item["metadata"]["solution"] - 1e-7)
+    assert dataset.score_answer(answer, item) == 1.0
+
+    # Answer is within 1e-1 of solution
+    answer = str(item["metadata"]["solution"] - 1e-2)
+    assert dataset.score_answer(answer, item) == 0.5
+
+    # Answer is far from solution
+    answer = str(item["metadata"]["solution"] - 1)
+    assert dataset.score_answer(answer, item) == 0.01
+
+    # Answer is None
+    answer = None
+    assert dataset.score_answer(answer, item) == 0.0