mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2026-04-19 12:58:07 +00:00
227 lines
8.6 KiB
Python
227 lines
8.6 KiB
Python
import cmath
|
||
import math
|
||
import random
|
||
from dataclasses import dataclass, field
|
||
from typing import Optional
|
||
|
||
from ..coaching import BaseCurriculum, ScalarAttributeDefinition
|
||
from ..factory import ProceduralDataset, register_dataset
|
||
|
||
DATASET_NAME = "complex_arithmetic"
|
||
|
||
|
||
@dataclass
|
||
class ComplexArithmeticConfig:
|
||
min_real: int = -10
|
||
max_real: int = 10
|
||
min_imag: int = -10
|
||
max_imag: int = 10
|
||
operations: tuple[str, ...] = ("+", "-", "*", "/")
|
||
operations_weights: list[float] = field(default_factory=lambda: [0.4, 0.4, 0.1, 0.1])
|
||
seed: Optional[int] = None
|
||
size: int = 500
|
||
|
||
def validate(self) -> None:
|
||
"""Validate configuration parameters."""
|
||
assert self.max_real >= self.min_real, "max_real must be >= min_real"
|
||
assert self.max_imag >= self.min_imag, "max_imag must be >= min_imag"
|
||
assert all(op in ("+", "-", "*", "/") for op in self.operations), "invalid operator"
|
||
assert round(sum(self.operations_weights), 1) == 1.0, "operations_weights must sum to 1.0"
|
||
|
||
|
||
class ComplexArithmeticDataset(ProceduralDataset):
|
||
"""Generates complex number arithmetic problems."""
|
||
|
||
def __init__(self, config: ComplexArithmeticConfig):
|
||
self._prompt_templates = {
|
||
"+": "Add the complex numbers: ({a}) + ({b})",
|
||
"-": "Subtract the complex numbers: ({a}) - ({b})",
|
||
"*": "Multiply the complex numbers: ({a}) × ({b})",
|
||
"/": "Divide the complex numbers: ({a}) ÷ ({b})",
|
||
}
|
||
super().__init__(config=config, seed=config.seed, size=config.size)
|
||
|
||
def _generate_complex(self, rng: random.Random) -> complex:
|
||
"""Generate a random complex number."""
|
||
real = rng.randint(self.config.min_real, self.config.max_real)
|
||
imag = rng.randint(self.config.min_imag, self.config.max_imag)
|
||
return complex(real, imag)
|
||
|
||
def _format_complex(self, z: complex) -> str:
|
||
"""Format complex number with 2 decimal places."""
|
||
real, imag = z.real, z.imag
|
||
if abs(imag) < 1e-10:
|
||
return f"{real:.2f}"
|
||
elif abs(real) < 1e-10:
|
||
return f"{imag:.2f}i"
|
||
else:
|
||
sign = "+" if imag >= 0 else "-"
|
||
return f"{real} {sign} {abs(imag)}i"
|
||
|
||
def __getitem__(self, idx: int) -> dict:
|
||
rng = random.Random(self.seed + idx)
|
||
|
||
# Choose random operation
|
||
op = rng.choices(self.config.operations, weights=self.config.operations_weights, k=1)[0]
|
||
|
||
if op == "/":
|
||
# For division, first generate the quotient (a) and divisor (b)
|
||
# Then calculate the dividend (result) as a * b
|
||
a = self._generate_complex(rng) # This will be the final result
|
||
b = self._generate_complex(rng)
|
||
while b == 0: # Ensure non-zero divisor
|
||
b = self._generate_complex(rng)
|
||
result = a # Store the intended result
|
||
a = result * b # Calculate dividend to ensure whole number division
|
||
else:
|
||
# For other operations, generate numbers normally
|
||
a = self._generate_complex(rng)
|
||
b = self._generate_complex(rng)
|
||
|
||
# Calculate result
|
||
if op == "+":
|
||
result = a + b
|
||
elif op == "-":
|
||
result = a - b
|
||
else: # op == "*"
|
||
result = a * b
|
||
|
||
question = self._prompt_templates[op].format(a=self._format_complex(a), b=self._format_complex(b))
|
||
|
||
return {
|
||
"question": question,
|
||
"answer": self._format_complex(result),
|
||
"metadata": {
|
||
"source_dataset": DATASET_NAME,
|
||
"source_index": idx,
|
||
"num1": (a.real, a.imag),
|
||
"num2": (b.real, b.imag),
|
||
"operation": op,
|
||
"result": (int(result.real), int(result.imag)), # Convert to int since we ensure whole numbers
|
||
"difficulty": {
|
||
"min_real": self.config.min_real,
|
||
"max_real": self.config.max_real,
|
||
"min_imag": self.config.min_imag,
|
||
"max_imag": self.config.max_imag,
|
||
"operations_weights": self.config.operations_weights,
|
||
},
|
||
},
|
||
}
|
||
|
||
@staticmethod
|
||
def parse_string_to_complex(answer: str) -> complex:
|
||
if answer is None:
|
||
return None
|
||
|
||
try:
|
||
# Normalize the answer string by removing spaces and converting to lowercase
|
||
answer = answer.replace(" ", "").lower()
|
||
|
||
# remove brackets
|
||
while len(answer) > 1 and answer[0] == "(" and answer[-1] == ")":
|
||
answer = answer[1:-1]
|
||
|
||
# Convert mathematical notation 'i' to Python's 'j' for complex numbers
|
||
answer = answer.replace("i", "j")
|
||
|
||
# Handle real numbers (no imaginary part)
|
||
if "j" not in answer:
|
||
return complex(float(answer))
|
||
|
||
# Handle pure imaginary numbers
|
||
if answer == "j":
|
||
return complex(0, 1)
|
||
if answer == "-j":
|
||
return complex(0, -1)
|
||
|
||
# Handle cases like "7j" or "-7j" (no real part)
|
||
if (
|
||
answer.endswith("j")
|
||
and not any(c in answer[:-1] for c in "+-")
|
||
or (answer.startswith("-") and not any(c in answer[1:-1] for c in "+-"))
|
||
):
|
||
# Extract coefficient
|
||
coef = answer[:-1]
|
||
if coef == "":
|
||
coef = "1"
|
||
elif coef == "-":
|
||
coef = "-1"
|
||
return complex(0, float(coef))
|
||
|
||
# Handle complex numbers with both parts
|
||
# Make sure there's a + or - before j if it's not at the beginning
|
||
if "j" in answer and not answer.endswith("j"):
|
||
return None # Invalid format like "3j+2"
|
||
|
||
# Handle cases like "3+j" (implicit 1)
|
||
if "+j" in answer:
|
||
answer = answer.replace("+j", "+1j")
|
||
if "-j" in answer:
|
||
answer = answer.replace("-j", "-1j")
|
||
|
||
# Parse the normalized string into a complex number
|
||
return complex(answer)
|
||
|
||
except (ValueError, TypeError):
|
||
return None
|
||
|
||
def score_answer(self, answer: Optional[str], entry: dict) -> float:
|
||
"""Score the answer using exponential distance-based scoring."""
|
||
if answer is None:
|
||
return 0.0
|
||
|
||
metadata = entry["metadata"]
|
||
try:
|
||
student_result = self.parse_string_to_complex(answer)
|
||
expected_result = complex(*metadata["result"])
|
||
# Calculate distance-based score using exponential decay
|
||
distance = abs(student_result - expected_result)
|
||
score = min(1.0, math.exp(-distance)) # Add 'import math' at the top
|
||
return score
|
||
|
||
except (ValueError, TypeError):
|
||
return 0.0
|
||
|
||
|
||
class ComplexArithmeticCurriculum(BaseCurriculum):
|
||
"""Curriculum for complex number arithmetic problems."""
|
||
|
||
def __init__(self):
|
||
super().__init__(ComplexArithmeticCurriculum.__name__, ComplexArithmeticConfig)
|
||
|
||
# Define attributes
|
||
self._define_attributes(
|
||
ScalarAttributeDefinition(
|
||
name="min_real",
|
||
field_name="min_real",
|
||
levels=[-10, -100, -10000, -100000000],
|
||
description="Minimum real part for complex numbers",
|
||
),
|
||
ScalarAttributeDefinition(
|
||
name="max_real",
|
||
field_name="max_real",
|
||
levels=[10, 100, 10000, 100000000],
|
||
description="Maximum real part for complex numbers",
|
||
),
|
||
ScalarAttributeDefinition(
|
||
name="min_imag",
|
||
field_name="min_imag",
|
||
levels=[-10, -100, -10000, -100000000],
|
||
description="Minimum imaginary part for complex numbers",
|
||
),
|
||
ScalarAttributeDefinition(
|
||
name="max_imag",
|
||
field_name="max_imag",
|
||
levels=[10, 100, 10000, 100000000],
|
||
description="Maximum imaginary part for complex numbers",
|
||
),
|
||
ScalarAttributeDefinition(
|
||
name="operations_weights",
|
||
field_name="operations_weights",
|
||
levels=[[0.4, 0.4, 0.1, 0.1], [0.25, 0.25, 0.25, 0.25], [0.2, 0.2, 0.3, 0.3], [0.1, 0.1, 0.4, 0.4]],
|
||
description="Operations weights to sample operation to use for each complex arithmetic problem",
|
||
),
|
||
)
|
||
|
||
|
||
register_dataset(DATASET_NAME, ComplexArithmeticDataset, ComplexArithmeticConfig, ComplexArithmeticCurriculum)
|