Power Function (#102)

* power function dataset + tests

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)