added impl of simple integration dataset generator

reasoning_gym/algebra/simple_integration.py

@@ -0,0 +1,80 @@
+import random
+from dataclasses import dataclass
+from fractions import Fraction
+from typing import Optional
+
+import sympy
+
+from ..factory import ProceduralDataset, register_dataset
+
+
+@dataclass
+class SimpleIntegrationConfig:
+    min_terms: int = 2
+    max_terms: int = 5
+    min_degree: int = 1
+    max_degree: int = 10
+    min_bounds: int = 1
+    max_bounds: int = 10
+    operators: tuple = ("+", "-")
+    symbols: tuple = ("x", "X")
+    seed: Optional[int] = None
+    size: int = 500
+
+    def validate(self) -> None:
+        """Validate the configuration parameters of the integral proble"""
+        assert self.min_bounds > 0, "min_bounds must be positive"
+        assert self.max_bounds >= self.min_bounds, "max_bounds must be >= min_bounds"
+        assert self.min_terms >= 0, "min_terms must be positive"
+        assert self.max_terms >= self.min_terms, "max_terms must be >= min_terms"
+        assert self.min_degree >= -10, "min_degree must be >= -10"
+        assert self.max_degree >= self.min_degree, "max_degree must be >= min_degree"
+        assert all(op in ("+", "-") for op in self.operators), "invalid operator specified"
+
+
+class SimpleIntegrationDataset(ProceduralDataset):
+    """Generates simple integration problems with one variable"""
+
+    def __init__(self, config: SimpleIntegrationConfig):
+        self._prompt_templates = [
+            "Find the indefinite integral: ∫ {integrand} dx",
+            "Calculate the antiderivative: ∫ {integrand} dx",
+            "Evaluate the indefinite integral: ∫ {integrand} dx",
+        ]
+        super().__init__(config=config, seed=config.seed, size=config.size)
+
+    def _generate_coefficient(self, rng: random.Random) -> Fraction:
+        """Generate a random coefficient for the polynomial"""
+        if rng.choice([True, False]):  # 50% chance for integer
+            return Fraction(rng.randint(self.config.min_bounds, self.config.max_bounds), 1)
+        denominator = rng.randint(2, 10)
+        return Fraction(rng.randint(self.config.min_bounds, self.config.max_bounds), denominator)
+
+    def _generate_polynomial(self, rng: random.Random) -> tuple[sympy.Symbol, sympy.Expr]:
+        """Generate a random polynomial with one variable"""
+        terms = []
+        x = sympy.Symbol(rng.choice(self.config.symbols))
+
+        for _ in range(rng.randint(self.config.min_terms, self.config.max_terms)):
+            coefficient = self._generate_coefficient(rng)
+            degree = rng.randint(self.config.min_degree, self.config.max_degree)
+            operator = rng.choice(self.config.operators)
+            term = coefficient * x**degree
+            if operator == "-":
+                term = -term
+            terms.append(term)
+        return x, sum(terms)
+
+    def __getitem__(self, idx: int) -> dict:
+        rng = random.Random(self.seed + idx)
+        symbol, polynomial = self._generate_polynomial(rng)
+        derivative = sympy.diff(polynomial, symbol)
+
+        return {
+            "question": rng.choice(self._prompt_templates).format(integrand=derivative),
+            "answer": str(polynomial) + " + C",
+            "metadata": {"integrand": str(derivative), "variable": str(symbol), "antiderivative": str(polynomial)},
+        }
+
+
+register_dataset("simple_integration", SimpleIntegrationDataset, SimpleIntegrationConfig)