Paterson-Stockmeyer algorithm in Python
Evaluates a degree-n polynomial with O(sqrt(n)) nonscalar multiplications by baby-step giant-step splitting.
- File
paterson_stockmeyer_algorithm.py, 185 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/paterson-stockmeyer-algorithm/. The authors are listed in prov/E2T01W.jsonl.
"""The Paterson-Stockmeyer algorithm for evaluating a polynomial at an algebra element.
A polynomial a_0 + a_1 x + ... + a_n x^n is the list [a_0, ..., a_n], lowest
degree first; the zero polynomial is []. The coefficients are scalars from a
commutative ring R (int, Fraction, a modular class), and the point X is an
element of an R-algebra A with + and *, for example a square matrix. A scalar
c times an element P of A is written scale(c, P), by default c * P.
A product of two elements of A is a nonscalar multiplication, and is the
expensive operation. The algorithm computes the baby steps X^0, ..., X^k,
forms the blocks B_j = sum_i a_(jk+i) X^i with scalar multiplications only,
and combines the blocks by Horner's rule in Y = X^k. With k about sqrt(n + 1)
it does at most 2 ceil(sqrt(n + 1)) - 2 nonscalar multiplications.
"""
from __future__ import annotations
from math import isqrt
from typing import Any, Callable, List, Optional, Sequence, Tuple
def _default_scale(c: Any, p: Any) -> Any:
return c * p
def block_size(n: int) -> int:
"""Return ceil(sqrt(n + 1)), the default block size for degree n."""
k = isqrt(n + 1)
return k if k * k == n + 1 else k + 1
def paterson_stockmeyer(
a: Sequence[Any],
x: Any,
one: Any = 1,
k: Optional[int] = None,
scale: Callable[[Any, Any], Any] = _default_scale,
) -> Any:
"""Return a_0 one + a_1 x + ... + a_n x^n.
one is the unit of the algebra (1 for scalars, the identity matrix for
matrices). k is the block size; the default is ceil(sqrt(n + 1)).
"""
a = list(a)
if not a:
return scale(0, one)
n = len(a) - 1
if k is None:
k = block_size(n) # @line L1
if k < 1:
raise ValueError("the block size must be at least 1")
m = -(-(n + 1) // k)
powers = [one, x] # @line L2
for _ in range(2, min(k, n) + 1): # @line L3
powers.append(powers[-1] * x)
blocks = []
for j in range(m): # @line L4
b = scale(a[j * k], one)
for i in range(1, k):
t = j * k + i
if t > n:
break
b = b + scale(a[t], powers[i])
blocks.append(b)
s = blocks[m - 1] # @line L5
for j in range(m - 2, -1, -1): # @line L6
s = s * powers[k] + blocks[j] # @line L7
return s # @line L8
def horner(a: Sequence[Any], x: Any, one: Any = 1, scale: Callable[[Any, Any], Any] = _default_scale) -> Any:
"""Return a(x) by Horner's rule, with n nonscalar multiplications."""
a = list(a)
if not a:
return scale(0, one)
s = scale(a[-1], one)
for c in reversed(a[:-1]):
s = s * x + scale(c, one)
return s
def nonscalar_count(n: int, k: Optional[int] = None) -> int:
"""Return the number of nonscalar products PATERSON-STOCKMEYER does for degree n."""
if n < 1:
return 0
if k is None:
k = block_size(n)
m = -(-(n + 1) // k)
return max(min(k, n) - 1, 0) + (m - 1)
class Counted:
"""An algebra element that counts nonscalar and scalar multiplications."""
stats = {"nonscalar": 0, "scalar": 0}
def __init__(self, v: Any) -> None:
self.v = v
def __add__(self, o: "Counted") -> "Counted":
return Counted(self.v + o.v)
def __mul__(self, o: "Counted") -> "Counted":
Counted.stats["nonscalar"] += 1
return Counted(self.v * o.v)
@staticmethod
def scale(c: Any, p: "Counted") -> "Counted":
Counted.stats["scalar"] += 1
return Counted(c * p.v)
def counted_run(a: Sequence[Any], x: Any, k: Optional[int] = None, method: str = "ps") -> Tuple[Any, int, int]:
"""Run a method on a scalar point and return (value, nonscalar, scalar)."""
Counted.stats = {"nonscalar": 0, "scalar": 0}
if method == "ps":
r = paterson_stockmeyer(a, Counted(x), Counted(1), k, Counted.scale)
else:
r = horner(a, Counted(x), Counted(1), Counted.scale)
return r.v, Counted.stats["nonscalar"], Counted.stats["scalar"]
def trace(a: Sequence[Any], x: Any, k: Optional[int] = None) -> Tuple[List[Any], List[Any], List[Any]]:
"""Return (powers, blocks, Horner partial values) for a scalar point."""
n = len(a) - 1
if k is None:
k = block_size(n)
m = -(-(n + 1) // k)
powers = [x**i for i in range(k + 1)]
blocks = [sum(a[t] * powers[t - j * k] for t in range(j * k, min(j * k + k, n + 1))) for j in range(m)]
s = blocks[-1]
partial = [s]
for j in range(m - 2, -1, -1):
s = s * powers[k] + blocks[j]
partial.append(s)
return powers, blocks, partial
class Mat:
"""A small square matrix over a commutative ring, for matrix arguments."""
def __init__(self, rows: Sequence[Sequence[Any]]) -> None:
self.r = [list(row) for row in rows]
@staticmethod
def identity(d: int, one: Any = 1, zero: Any = 0) -> "Mat":
return Mat([[one if i == j else zero for j in range(d)] for i in range(d)])
def __add__(self, o: "Mat") -> "Mat":
return Mat([[x + y for x, y in zip(r1, r2)] for r1, r2 in zip(self.r, o.r)])
def __mul__(self, o: "Mat") -> "Mat":
d = len(self.r)
cols = list(zip(*o.r))
out = []
for i in range(d):
row = []
for j in range(d):
s = self.r[i][0] * cols[j][0]
for t in range(1, d):
s = s + self.r[i][t] * cols[j][t]
row.append(s)
out.append(row)
return Mat(out)
def __eq__(self, o: object) -> bool:
return isinstance(o, Mat) and self.r == o.r
def __repr__(self) -> str:
return f"Mat({self.r})"
@staticmethod
def scale(c: Any, p: "Mat") -> "Mat":
return Mat([[c * x for x in row] for row in p.r])
if __name__ == "__main__":
a = [1, 2, 3, 4, 5, 6, 7, 8, 9]
print(trace(a, 2))
print(counted_run(a, 2), counted_run(a, 2, method="horner"))
j = Mat([[1, 1], [0, 1]])
print(paterson_stockmeyer(a, j, Mat.identity(2), scale=Mat.scale))
test_paterson_stockmeyer_algorithm.py 200 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/paterson-stockmeyer-algorithm/. The authors are listed in prov/E2T01W.jsonl.
import random
from fractions import Fraction
import pytest
from paterson_stockmeyer_algorithm import (
Mat,
block_size,
counted_run,
horner,
nonscalar_count,
paterson_stockmeyer,
trace,
)
class ModP:
"""An element of Z/pZ, enough for the tests."""
def __init__(self, v: int, p: int) -> None:
self.p = p
self.v = v % p
def _c(self, o):
return o.v if isinstance(o, ModP) else o
def __add__(self, o):
return ModP(self.v + self._c(o), self.p)
__radd__ = __add__
def __sub__(self, o):
return ModP(self.v - self._c(o), self.p)
def __rsub__(self, o):
return ModP(self._c(o) - self.v, self.p)
def __mul__(self, o):
return ModP(self.v * self._c(o), self.p)
__rmul__ = __mul__
def __neg__(self):
return ModP(-self.v, self.p)
def __truediv__(self, o):
d = self._c(o) % self.p
if d == 0:
raise ZeroDivisionError
return ModP(self.v * pow(d, -1, self.p), self.p)
def __rtruediv__(self, o):
return ModP(o, self.p) / self
def __eq__(self, o):
return self.v == self._c(o) % self.p
def __hash__(self):
return hash(self.v)
def __repr__(self):
return f"{self.v} mod {self.p}"
# Independent oracles: the sum of a_i x^i with powers formed one by one.
def oracle_scalar(a, x, p=None):
if p:
return sum(c * pow(x, i, p) for i, c in enumerate(a)) % p
return sum(c * x**i for i, c in enumerate(a))
def mat_mul(u, v):
d = len(u)
return [[sum(u[i][t] * v[t][j] for t in range(d)) for j in range(d)] for i in range(d)]
def oracle_matrix(a, x, zero, one):
d = len(x)
acc = [[zero for _ in range(d)] for _ in range(d)]
pw = [[one if i == j else zero for j in range(d)] for i in range(d)]
for c in a:
acc = [[acc[i][j] + c * pw[i][j] for j in range(d)] for i in range(d)]
pw = mat_mul(pw, x)
return acc
PRIMES = [2, 3, 7, 998244353]
def test_page_example():
a = [1, 2, 3, 4, 5, 6, 7, 8, 9]
assert trace(a, 2) == ([1, 2, 4, 8], [17, 38, 59], [59, 510, 4097])
assert counted_run(a, 2) == (4097, 4, 9)
assert counted_run(a, 2, method="horner")[:2] == (4097, 8)
j = Mat([[1, 1], [0, 1]])
assert paterson_stockmeyer(a, j, Mat.identity(2), scale=Mat.scale) == Mat([[45, 240], [0, 45]])
def test_edge_cases():
assert paterson_stockmeyer([], 5) == 0
assert paterson_stockmeyer([7], 5) == 7
assert paterson_stockmeyer([7, 3], 5) == 22
assert paterson_stockmeyer([0, 0, 0], 5) == 0
# leading zero coefficients do not change the value
assert paterson_stockmeyer([1, 2, 0, 0, 0], 3) == 7
assert paterson_stockmeyer([2, 0, 0, 1], 0) == 2
with pytest.raises(ValueError):
paterson_stockmeyer([1, 2], 3, k=0)
z = Mat([[0, 0], [0, 0]])
assert paterson_stockmeyer([], z, Mat.identity(2), scale=Mat.scale) == z
assert paterson_stockmeyer([4], z, Mat.identity(2), scale=Mat.scale) == Mat([[4, 0], [0, 4]])
def test_over_z_random_and_all_block_sizes():
rng = random.Random(1)
for _ in range(300):
n = rng.randint(0, 60)
a = [rng.randint(-99, 99) for _ in range(n + 1)]
x = rng.randint(-5, 5)
want = oracle_scalar(a, x)
assert paterson_stockmeyer(a, x) == want
assert horner(a, x) == want
k = rng.randint(1, n + 3)
assert paterson_stockmeyer(a, x, k=k) == want
@pytest.mark.parametrize("p", PRIMES)
def test_over_z_mod_p(p):
rng = random.Random(p)
for _ in range(200):
n = rng.randint(0, 70)
a = [rng.randrange(p) for _ in range(n + 1)]
x = rng.randrange(p)
got = paterson_stockmeyer([ModP(c, p) for c in a], ModP(x, p), ModP(1, p))
assert got == oracle_scalar(a, x, p)
def test_over_fractions():
rng = random.Random(5)
for _ in range(150):
n = rng.randint(0, 30)
a = [Fraction(rng.randint(-20, 20), rng.randint(1, 9)) for _ in range(n + 1)]
x = Fraction(rng.randint(-7, 7), rng.randint(1, 5))
assert paterson_stockmeyer(a, x, Fraction(1)) == oracle_scalar(a, x)
def test_matrix_argument_over_z_and_mod_p():
rng = random.Random(11)
for _ in range(60):
d = rng.randint(1, 4)
n = rng.randint(0, 25)
a = [rng.randint(-9, 9) for _ in range(n + 1)]
x = [[rng.randint(-2, 2) for _ in range(d)] for _ in range(d)]
got = paterson_stockmeyer(a, Mat(x), Mat.identity(d), scale=Mat.scale)
assert got.r == oracle_matrix(a, x, 0, 1)
for p in PRIMES:
for _ in range(15):
d = rng.randint(1, 3)
n = rng.randint(0, 25)
a = [ModP(rng.randrange(p), p) for _ in range(n + 1)]
x = [[ModP(rng.randrange(p), p) for _ in range(d)] for _ in range(d)]
one, zero = ModP(1, p), ModP(0, p)
got = paterson_stockmeyer(a, Mat(x), Mat.identity(d, one, zero), scale=Mat.scale)
assert got.r == oracle_matrix(a, x, zero, one)
def test_matrix_over_fractions():
rng = random.Random(13)
for _ in range(20):
n = rng.randint(0, 15)
a = [Fraction(rng.randint(-9, 9), rng.randint(1, 4)) for _ in range(n + 1)]
x = [[Fraction(rng.randint(-3, 3), rng.randint(1, 3)) for _ in range(2)] for _ in range(2)]
one, zero = Fraction(1), Fraction(0)
got = paterson_stockmeyer(a, Mat(x), Mat.identity(2, one, zero), scale=Mat.scale)
assert got.r == oracle_matrix(a, x, zero, one)
def test_nonscalar_counts():
for n in range(0, 400):
k = block_size(n)
assert (k - 1) ** 2 < n + 1 <= k * k
a = list(range(1, n + 2))
v, ns, sc = counted_run(a, 3)
assert v == oracle_scalar(a, 3)
assert ns == nonscalar_count(n)
if n >= 1:
assert ns <= 2 * k - 2
assert sc <= n + 1 + k
for n in range(1, 60):
for k in range(1, n + 3):
a = [1] * (n + 1)
assert counted_run(a, 2, k=k)[1] == nonscalar_count(n, k)
# Horner uses n nonscalar products
assert counted_run(list(range(101)), 2, method="horner")[1] == 100
assert counted_run(list(range(101)), 2)[1] == 10 + 9
What the next level needs
- the tests did not run in this build