Estrin's scheme in Python
Evaluates a polynomial by a balanced tree of pairwise combinations, shortening the critical path on parallel hardware.
- File
estrin_scheme.py, 114 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/estrin-scheme/. The authors are listed in prov/SDH9BZ.jsonl.
"""Estrin's scheme for evaluating a polynomial with a short critical path.
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 and the point may be
int, Fraction, float, or any ring element type with + and *.
Each round combines neighbouring pairs c[2i] + c[2i+1] * p, where p is
x^(2^k) in round k, and then squares p. The products of one round do not
depend on each other, so on parallel hardware the critical path has about
2 log2(n + 1) operations, against 2n for Horner's rule.
"""
from __future__ import annotations
from typing import Any, List, Sequence, Tuple
def estrin(a: Sequence[Any], x: Any) -> Any:
"""Return a(x) by Estrin's scheme, with n + ceil(log2(n + 1)) - 1 products."""
c = list(a) # @line L1
if not c:
return x - x
p = x # @line L2
while len(c) > 1: # @line L3
nxt = [c[2 * i] + c[2 * i + 1] * p for i in range(len(c) // 2)] # @line L4
if len(c) % 2 == 1: # @line L5
nxt.append(c[-1])
c = nxt # @line L6
if len(c) > 1: # @line L7
p = p * p
return c[0] # @line L8
class Counter:
"""A ring element that counts operations and the depth of its expression."""
stats = {"mul": 0, "add": 0}
def __init__(self, v: Any, depth: int = 0) -> None:
self.v = v
self.depth = depth
def _w(self, o: Any) -> "Counter":
return o if isinstance(o, Counter) else Counter(o)
def __add__(self, o: Any) -> "Counter":
o = self._w(o)
Counter.stats["add"] += 1
return Counter(self.v + o.v, max(self.depth, o.depth) + 1)
__radd__ = __add__
def __sub__(self, o: Any) -> "Counter":
o = self._w(o)
return Counter(self.v - o.v, max(self.depth, o.depth))
def __mul__(self, o: Any) -> "Counter":
o = self._w(o)
Counter.stats["mul"] += 1
return Counter(self.v * o.v, max(self.depth, o.depth) + 1)
__rmul__ = __mul__
def estrin_stats(a: Sequence[Any], x: Any) -> Tuple[Any, int, int, int]:
"""Return (value, products, additions, depth) of Estrin's scheme."""
Counter.stats = {"mul": 0, "add": 0}
r = estrin([Counter(c) for c in a], Counter(x))
return r.v, Counter.stats["mul"], Counter.stats["add"], r.depth
def horner(a: Sequence[Any], x: Any) -> Any:
"""Reference: Horner's rule, n products and n additions in sequence."""
acc = x - x
for c in reversed(list(a)):
acc = acc * x + c
return acc
def horner_stats(a: Sequence[Any], x: Any) -> Tuple[Any, int, int, int]:
"""Return (value, products, additions, depth) of Horner's rule."""
Counter.stats = {"mul": 0, "add": 0}
a = list(a)
if not a:
return x - x, 0, 0, 0
acc = Counter(a[-1])
xc = Counter(x)
for c in reversed(a[:-1]):
acc = acc * xc + c
return acc.v, Counter.stats["mul"], Counter.stats["add"], acc.depth
def levels(a: Sequence[Any], x: Any) -> List[List[Any]]:
"""The lists c after each round, for traces."""
c = list(a)
out = [c]
p = x
while len(c) > 1:
nxt = [c[2 * i] + c[2 * i + 1] * p for i in range(len(c) // 2)]
if len(c) % 2 == 1:
nxt.append(c[-1])
c = nxt
out.append(c)
if len(c) > 1:
p = p * p
return out
if __name__ == "__main__":
print(levels([1, 2, 3, 4, 5, 6, 7, 8], 2))
print(estrin_stats([1, 2, 3, 4, 5, 6, 7, 8], 2), horner_stats([1, 2, 3, 4, 5, 6, 7, 8], 2))
test_estrin_scheme.py 140 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/estrin-scheme/. The authors are listed in prov/SDH9BZ.jsonl.
import math
import random
from fractions import Fraction
import pytest
from estrin_scheme import estrin, estrin_stats, horner, horner_stats, levels
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}"
def oracle(a, x, p=None):
if p:
return sum(c * pow(x, k, p) for k, c in enumerate(a)) % p
return sum(c * x**k for k, c in enumerate(a))
PRIMES = [2, 3, 7, 998244353]
def test_page_example():
assert levels([1, 2, 3, 4, 5, 6, 7, 8], 2) == [[1, 2, 3, 4, 5, 6, 7, 8], [5, 11, 17, 23], [49, 109], [1793]]
assert estrin_stats([1, 2, 3, 4, 5, 6, 7, 8], 2) == (1793, 9, 7, 6)
assert horner_stats([1, 2, 3, 4, 5, 6, 7, 8], 2) == (1793, 7, 7, 14)
def test_over_z_random():
rng = random.Random(1)
for _ in range(500):
n = rng.randint(0, 80)
a = [rng.randint(-99, 99) for _ in range(n)]
x = rng.randint(-20, 20)
assert estrin(a, x) == oracle(a, x) == horner(a, x)
@pytest.mark.parametrize("p", PRIMES)
def test_mod_p_random(p):
rng = random.Random(p)
for _ in range(300):
n = rng.randint(0, 60)
a = [rng.randrange(p) for _ in range(n)]
x = rng.randrange(p)
got = estrin([ModP(c, p) for c in a], ModP(x, p))
assert got == oracle(a, x, p)
def test_fraction_random():
rng = random.Random(5)
for _ in range(200):
n = rng.randint(0, 30)
a = [Fraction(rng.randint(-9, 9), rng.randint(1, 9)) for _ in range(n)]
x = Fraction(rng.randint(-9, 9), rng.randint(1, 9))
assert estrin(a, x) == oracle(a, x)
def test_operation_counts_and_depth():
for n in range(0, 130):
a = list(range(1, n + 2))
v, muls, adds, depth = estrin_stats(a, 3)
lg = math.ceil(math.log2(n + 1)) if n >= 1 else 0
assert v == oracle(a, 3)
assert adds == n
assert muls == n + max(lg - 1, 0)
assert depth <= 2 * lg
hv, hm, ha, hd = horner_stats(a, 3)
assert (hv, hm, ha, hd) == (v, n, n, 2 * n)
def test_floats_close_to_horner():
rng = random.Random(9)
for _ in range(100):
a = [rng.uniform(-1, 1) for _ in range(rng.randint(1, 40))]
x = rng.uniform(-1, 1)
assert math.isclose(estrin(a, x), horner(a, x), rel_tol=1e-9, abs_tol=1e-12)
def test_edge_cases():
assert estrin([], 5) == 0
assert estrin([0, 0, 0], 5) == 0
assert estrin([7], 5) == 7
assert estrin([3, 2], 5) == 13
# Leading zero coefficients do not change the value.
assert estrin([1, 2, 3, 0, 0], 4) == estrin([1, 2, 3], 4) == 57
# Point zero and point one.
assert estrin([4, 5, 6, 7, 8], 0) == 4
assert estrin([4, 5, 6, 7, 8], 1) == 30
# Odd lengths, so a coefficient is carried up.
assert levels([1, 1, 1], 10) == [[1, 1, 1], [11, 1], [111]]
assert estrin([1] * 7, 10) == 1111111
# Over Z/2.
assert estrin([ModP(1, 2)] * 5, ModP(1, 2)) == 1
What the next level needs
- the tests did not run in this build