diff --git a/pyproject.toml b/pyproject.toml index 4b852c2..61e6732 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,3 +15,6 @@ ignore = [ # test_smoke.py star-imports sage.all and loads a .sage file, so ruff can't see where its names come from. "tests/test_topdata_and_disc.py" = ["F403", "F405"] + +# Same star-import + load() pattern as above. +"tests/test_period_computation.py" = ["F403", "F405"] diff --git a/tests/test_period_computation.py b/tests/test_period_computation.py new file mode 100644 index 0000000..4c85e7d --- /dev/null +++ b/tests/test_period_computation.py @@ -0,0 +1,48 @@ +import os + +import pytest +from sage.all import * # noqa: F401 + +load(os.path.join(os.path.dirname(__file__), "..", "sage", "period_computation.sage")) + + +def test_apply_operator_univariate(): + # theta^2 - z applied to log(z) should give -z*log(z), since theta(log z) = 1 + # and theta^2(log z) = theta(1) = 0. + period = Period(no_variables=1, coefficients={(1,): {(0,): 1}}, order=5) + op = PFOperator("theta0^2 - z0", no_variables=1) + + result = period.apply_operator(op) + + assert result.coefficients == {(1,): {(1,): -1}} + + +def test_apply_operator_mixes_variables(): + # z0*theta1 applied to z0*log(z1): theta1 strips log(z1) down to a bare 1 + # (leaving z0 untouched), then multiplying by z0 gives z0^2. + period = Period(no_variables=2, coefficients={(0, 1): {(1, 0): 1}}, order=5) + op = PFOperator("z0*theta1", no_variables=2) + + result = period.apply_operator(op) + + assert result.coefficients == {(0, 0): {(2, 0): 1}} + + +def test_apply_operator_truncates_to_order(): + # Multiplying by z0^2 pushes some terms above the period's order, so they + # should be dropped rather than kept with a nonzero coefficient. + period = Period(no_variables=1, coefficients={(0,): {(0,): 1, (1,): 1, (2,): 1}}, order=2) + op = PFOperator("z0^2", no_variables=1) + + result = period.apply_operator(op) + + assert result.coefficients == {(0,): {(2,): 1}} + assert result.order == 2 + + +def test_apply_operator_rejects_variable_count_mismatch(): + period = Period(no_variables=1, coefficients={}) + op = PFOperator("z0*theta1", no_variables=2) + + with pytest.raises(ValueError): + period.apply_operator(op)