mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2026-04-19 12:58:07 +00:00
201 lines
8.5 KiB
Python
201 lines
8.5 KiB
Python
import pytest
|
|
|
|
from reasoning_gym.algebra.complex_arithmetic import (
|
|
ComplexArithmeticConfig,
|
|
ComplexArithmeticCurriculum,
|
|
ComplexArithmeticDataset,
|
|
)
|
|
|
|
|
|
def test_complex_arithmetic_basic():
|
|
"""Test basic functionality of complex arithmetic dataset."""
|
|
config = ComplexArithmeticConfig(
|
|
min_real=-5, max_real=5, min_imag=-5, max_imag=5, operations=("+", "-", "*", "/"), seed=42, size=10
|
|
)
|
|
dataset = ComplexArithmeticDataset(config)
|
|
|
|
print(dataset)
|
|
|
|
# Test dataset size
|
|
assert len(dataset) == 10
|
|
|
|
# Test a specific item
|
|
item = dataset[0]
|
|
assert "question" in item
|
|
assert "answer" in item
|
|
assert "metadata" in item
|
|
|
|
# Add more detailed assertions
|
|
assert isinstance(item["question"], str)
|
|
assert isinstance(item["answer"], str)
|
|
assert isinstance(item["metadata"], dict)
|
|
|
|
# Check metadata structure
|
|
assert "num1" in item["metadata"]
|
|
assert "num2" in item["metadata"]
|
|
assert "operation" in item["metadata"]
|
|
assert "result" in item["metadata"]
|
|
|
|
# Check data types in metadata
|
|
assert isinstance(item["metadata"]["num1"], tuple)
|
|
assert isinstance(item["metadata"]["num2"], tuple)
|
|
assert len(item["metadata"]["num1"]) == 2 # Real and imaginary parts
|
|
assert len(item["metadata"]["num2"]) == 2
|
|
assert isinstance(item["metadata"]["operation"], str)
|
|
assert isinstance(item["metadata"]["result"], tuple)
|
|
|
|
# Make sure answer matches the result in metadata
|
|
# results is a tuple of two floats (real, imag) and answer is a string
|
|
# answer is formatted as "real + imagi"
|
|
assert ComplexArithmeticDataset.parse_string_to_complex(item["answer"]) == complex(*item["metadata"]["result"])
|
|
|
|
|
|
def test_complex_arithmetic_scoring():
|
|
"""Test scoring function with various answer formats and accuracies."""
|
|
config = ComplexArithmeticConfig(seed=42)
|
|
dataset = ComplexArithmeticDataset(config)
|
|
|
|
# Test case with answer 3 + 2i
|
|
entry = {"metadata": {"result": (3.0, 2.0)}}
|
|
|
|
# Test exact matches (should get score of 1.0)
|
|
assert dataset.score_answer("3 + 2i", entry) == 1.0
|
|
assert dataset.score_answer("3+2i", entry) == 1.0
|
|
assert dataset.score_answer("3.0 + 2.0i", entry) == 1.0
|
|
assert dataset.score_answer("((3.0 + 2.0i ) )", entry) == 1.0
|
|
|
|
# Test answers with small errors (should get high but < 1.0 scores)
|
|
print(dataset.score_answer("3.1 + 2i", entry))
|
|
assert 0.9 < dataset.score_answer("3.1 + 2i", entry) < 1.0
|
|
assert 0.9 < dataset.score_answer("3 + 2.1i", entry) < 1.0
|
|
assert 0.7 < dataset.score_answer("3.1 + 2.1i", entry) < 0.95
|
|
|
|
# Test answers with moderate errors (should get medium scores)
|
|
assert 0.3 < dataset.score_answer("4 + 2i", entry) < 0.4
|
|
assert 0.3 < dataset.score_answer("3 + 3i", entry) < 0.4
|
|
|
|
# Test answers with large errors (should get very low scores)
|
|
assert dataset.score_answer("10 + 10i", entry) < 0.01
|
|
|
|
# Test invalid answers (should get 0.0)
|
|
assert dataset.score_answer("invalid", entry) == 0.0
|
|
assert dataset.score_answer(None, entry) == 0.0
|
|
assert dataset.score_answer("inf + 2i", entry) == 0.0
|
|
|
|
|
|
def test_complex_arithmetic_division_by_zero():
|
|
"""Test that division by zero is handled properly."""
|
|
config = ComplexArithmeticConfig(
|
|
operations=("+", "-", "*", "/"), operations_weights=[0.0, 0.0, 0.0, 1.0], seed=42
|
|
) # Only test division
|
|
dataset = ComplexArithmeticDataset(config)
|
|
|
|
# Check multiple items to ensure no division by zero
|
|
for i in range(10):
|
|
item = dataset[i]
|
|
num2 = complex(*item["metadata"]["num2"])
|
|
assert num2 != 0
|
|
|
|
|
|
def test_parse_string_to_complex():
|
|
"""Test the parse_string_to_complex method with various input formats."""
|
|
dataset = ComplexArithmeticDataset(ComplexArithmeticConfig())
|
|
|
|
# Test pure real numbers
|
|
assert dataset.parse_string_to_complex("5") == complex(5, 0)
|
|
assert dataset.parse_string_to_complex("5.0") == complex(5, 0)
|
|
assert dataset.parse_string_to_complex("-5") == complex(-5, 0)
|
|
assert dataset.parse_string_to_complex("-5.5") == complex(-5.5, 0)
|
|
|
|
# Test pure imaginary numbers
|
|
assert dataset.parse_string_to_complex("i") == complex(0, 1)
|
|
assert dataset.parse_string_to_complex("j") == complex(0, 1)
|
|
assert dataset.parse_string_to_complex("7i") == complex(0, 7)
|
|
assert dataset.parse_string_to_complex("7.0i") == complex(0, 7)
|
|
assert dataset.parse_string_to_complex("-i") == complex(0, -1)
|
|
assert dataset.parse_string_to_complex("-7i") == complex(0, -7)
|
|
assert dataset.parse_string_to_complex("-7.5i") == complex(0, -7.5)
|
|
|
|
# Test complex numbers with both parts
|
|
assert dataset.parse_string_to_complex("3+2i") == complex(3, 2)
|
|
assert dataset.parse_string_to_complex("3 + 2i") == complex(3, 2)
|
|
assert dataset.parse_string_to_complex("3.5 + 2.5i") == complex(3.5, 2.5)
|
|
assert dataset.parse_string_to_complex("-3 + 2i") == complex(-3, 2)
|
|
assert dataset.parse_string_to_complex("3 - 2i") == complex(3, -2)
|
|
assert dataset.parse_string_to_complex("-3 - 2i") == complex(-3, -2)
|
|
|
|
# Test with j notation
|
|
assert dataset.parse_string_to_complex("3+2j") == complex(3, 2)
|
|
assert dataset.parse_string_to_complex("3 + 2j") == complex(3, 2)
|
|
|
|
# Test with different spacing
|
|
assert dataset.parse_string_to_complex(" 3 + 2i ") == complex(3, 2)
|
|
assert dataset.parse_string_to_complex("3+i") == complex(3, 1)
|
|
assert dataset.parse_string_to_complex("3-i") == complex(3, -1)
|
|
|
|
# Test invalid inputs
|
|
assert dataset.parse_string_to_complex("invalid") is None
|
|
assert dataset.parse_string_to_complex("3 + i + 2") is None
|
|
assert dataset.parse_string_to_complex("3 + 2x") is None
|
|
|
|
|
|
def test_complex_arithmetic_curriculum():
|
|
"""Test the curriculum for complex arithmetic."""
|
|
curriculum = ComplexArithmeticCurriculum()
|
|
base_value = {"size": 150, "seed": 1}
|
|
|
|
base_cfg: ComplexArithmeticCurriculum = curriculum.generate_configuration(base_value)
|
|
|
|
assert base_cfg.seed == 1
|
|
assert base_cfg.size == 150
|
|
assert base_cfg.min_real == base_cfg.min_imag == -10
|
|
assert base_cfg.max_real == base_cfg.max_imag == 10
|
|
assert base_cfg.operations_weights == [0.4, 0.4, 0.1, 0.1]
|
|
|
|
# Increase and validate increase in level
|
|
curriculum.increment_attr_level("min_real")
|
|
curriculum.increment_attr_level("min_imag")
|
|
curriculum.increment_attr_level("max_real")
|
|
curriculum.increment_attr_level("max_imag")
|
|
curriculum.increment_attr_level("operations_weights")
|
|
|
|
increased_cfg: ComplexArithmeticCurriculum = curriculum.generate_configuration(base_value)
|
|
assert increased_cfg.min_real == increased_cfg.min_imag == -100
|
|
assert increased_cfg.max_real == increased_cfg.max_imag == 100
|
|
assert increased_cfg.operations_weights == [0.25, 0.25, 0.25, 0.25]
|
|
|
|
# Decrease and validate decrease in level
|
|
curriculum.decrement_attr_level("min_real")
|
|
curriculum.decrement_attr_level("min_imag")
|
|
curriculum.decrement_attr_level("max_real")
|
|
curriculum.decrement_attr_level("max_imag")
|
|
curriculum.decrement_attr_level("operations_weights")
|
|
|
|
decreased_cfg: ComplexArithmeticCurriculum = curriculum.generate_configuration(base_value)
|
|
assert decreased_cfg.min_real == decreased_cfg.min_imag == -10
|
|
assert decreased_cfg.max_real == decreased_cfg.max_imag == 10
|
|
assert decreased_cfg.operations_weights == [0.4, 0.4, 0.1, 0.1]
|
|
|
|
# Test upper bound boundary condition
|
|
for _ in range(10):
|
|
curriculum.increment_attr_level("min_real")
|
|
curriculum.increment_attr_level("min_imag")
|
|
curriculum.increment_attr_level("max_real")
|
|
curriculum.increment_attr_level("max_imag")
|
|
curriculum.increment_attr_level("operations_weights")
|
|
upper_bound_cfg: ComplexArithmeticCurriculum = curriculum.generate_configuration(base_value)
|
|
assert upper_bound_cfg.min_real == upper_bound_cfg.min_imag == -100000000
|
|
assert upper_bound_cfg.max_real == upper_bound_cfg.max_imag == 100000000
|
|
assert upper_bound_cfg.operations_weights == [0.1, 0.1, 0.4, 0.4]
|
|
|
|
# Test lower bound boundary condition
|
|
for _ in range(10):
|
|
curriculum.decrement_attr_level("min_real")
|
|
curriculum.decrement_attr_level("min_imag")
|
|
curriculum.decrement_attr_level("max_real")
|
|
curriculum.decrement_attr_level("max_imag")
|
|
curriculum.decrement_attr_level("operations_weights")
|
|
lower_bound_cfg: ComplexArithmeticCurriculum = curriculum.generate_configuration(base_value)
|
|
assert lower_bound_cfg.min_real == lower_bound_cfg.min_imag == -10
|
|
assert lower_bound_cfg.max_real == lower_bound_cfg.max_imag == 10
|
|
assert lower_bound_cfg.operations_weights == [0.4, 0.4, 0.1, 0.1]
|