mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2026-04-19 12:58:07 +00:00
117 lines
3.6 KiB
Python
117 lines
3.6 KiB
Python
import pytest
|
|
from sympy import Symbol, sympify
|
|
|
|
from reasoning_gym import create_dataset
|
|
from reasoning_gym.algebra.polynomial_equations import PolynomialEquationsConfig, PolynomialEquationsDataset
|
|
|
|
|
|
def test_polynomial_config_validation():
|
|
"""Test that invalid configs raise appropriate errors"""
|
|
with pytest.raises(AssertionError):
|
|
PolynomialEquationsConfig(min_terms=0).validate()
|
|
|
|
with pytest.raises(AssertionError):
|
|
PolynomialEquationsConfig(min_value=0).validate()
|
|
|
|
with pytest.raises(AssertionError):
|
|
PolynomialEquationsConfig(min_degree=0, max_degree=3).validate()
|
|
|
|
with pytest.raises(AssertionError):
|
|
PolynomialEquationsConfig(min_degree=4, max_degree=3).validate()
|
|
|
|
with pytest.raises(AssertionError):
|
|
PolynomialEquationsConfig(operators=("^",)).validate()
|
|
|
|
|
|
def test_polynomial_equations_dataset_basic():
|
|
"""Test dataset creation and length"""
|
|
dataset_size = 50
|
|
config = PolynomialEquationsConfig(
|
|
min_terms=2,
|
|
max_terms=3,
|
|
min_value=1,
|
|
max_value=5,
|
|
min_degree=1,
|
|
max_degree=2,
|
|
seed=42,
|
|
size=dataset_size,
|
|
)
|
|
|
|
dataset = PolynomialEquationsDataset(config)
|
|
|
|
assert len(dataset) == dataset_size
|
|
|
|
|
|
def test_polynomial_equations_dataset_items():
|
|
"""Test that generated items have correct structure"""
|
|
ds = create_dataset(
|
|
"polynomial_equations",
|
|
min_terms=2,
|
|
max_terms=3,
|
|
min_value=1,
|
|
max_value=5,
|
|
min_degree=1,
|
|
max_degree=2,
|
|
size=3,
|
|
seed=100,
|
|
)
|
|
|
|
for item in ds:
|
|
assert "question" in item
|
|
assert "answer" in item
|
|
assert "metadata" in item
|
|
|
|
# Check metadata
|
|
assert isinstance(item["metadata"]["polynomial_expr"], str)
|
|
assert isinstance(item["metadata"]["variable"], str)
|
|
assert isinstance(item["metadata"]["degree"], int)
|
|
assert isinstance(item["metadata"]["real_solutions"], list)
|
|
|
|
# Check polynomial_expr existence
|
|
poly_str = item["metadata"]["polynomial_expr"]
|
|
# Ensure it can parse with sympy
|
|
sympify(poly_str)
|
|
|
|
|
|
def test_polynomial_equations_dataset_deterministic():
|
|
"""Test dataset reproducibility with fixed seed."""
|
|
cfg = PolynomialEquationsConfig(seed=999, size=3)
|
|
ds1 = PolynomialEquationsDataset(cfg)
|
|
ds2 = PolynomialEquationsDataset(cfg)
|
|
|
|
for i in range(len(ds1)):
|
|
assert ds1[i] == ds2[i], "Polynomial datasets with same seed should match exactly."
|
|
|
|
|
|
def test_polynomial_solutions_evaluation():
|
|
"""Test that real_solutions satisfy the polynomial equation."""
|
|
ds = create_dataset(
|
|
"polynomial_equations",
|
|
min_terms=2,
|
|
max_terms=4,
|
|
min_value=1,
|
|
max_value=10,
|
|
min_degree=1,
|
|
max_degree=3,
|
|
size=5,
|
|
seed=42,
|
|
)
|
|
|
|
for item in ds:
|
|
# Extract the polynomial expression and solutions
|
|
poly_str = item["metadata"]["polynomial_expr"]
|
|
real_solutions = item["metadata"]["real_solutions"]
|
|
x = Symbol(item["metadata"]["variable"])
|
|
# Parse the polynomial expression
|
|
poly_expr = sympify(poly_str)
|
|
|
|
# Verify that each solution satisfies the polynomial
|
|
for solution in real_solutions:
|
|
# Evaluate the expression with the solution substituted
|
|
evaluated_value = poly_expr.subs(x, solution)
|
|
|
|
# Ensure the evaluated value is close to zero (numerical stability threshold)
|
|
assert abs(evaluated_value) < 1e-6, (
|
|
f"Solution {solution} does not satisfy the polynomial {poly_str}. "
|
|
f"Evaluated value: {evaluated_value}"
|
|
)
|