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"]

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)