diff --git a/reasoning_gym/algebra/simple_integration.py b/reasoning_gym/algebra/simple_integration.py new file mode 100644 index 00000000..e524e3ef --- /dev/null +++ b/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)