Newton power series inversion in Python
Inverts a power series by the iteration g <- g(2 - fg) that doubles the precision each step.
- File
newton_power_series_inversion.py, 173 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/newton-power-series-inversion/. The authors are listed in prov/8P4715.jsonl.
"""Inversion of a power series by Newton iteration.
A power series f = f_0 + f_1 x + ... is given by a list [f_0, f_1, ...] of its
first coefficients, lowest degree first; missing coefficients are zero. The
coefficients may be int, Fraction, or any ring element type with + - * and
`/` for the inverse of the constant term (for example a Z/pZ class). The
result of inverse(f, n) is the list [g_0, ..., g_{n-1}] with f g = 1 mod x^n.
Costs are counted in ring operations. The multiplication is a parameter;
the default is Karatsuba multiplication, which works over every commutative
ring and has recursion depth O(log n).
"""
from __future__ import annotations
from typing import Any, Callable, List, Sequence
Poly = List[Any]
MulFn = Callable[[Sequence[Any], Sequence[Any]], Poly]
KARATSUBA_CUTOFF = 32
def schoolbook_mul(a: Sequence[Any], b: Sequence[Any]) -> Poly:
"""Return the full product of two coefficient lists, len(a) + len(b) - 1 entries."""
if not a or not b:
return []
c: Poly = [a[0] - a[0]] * (len(a) + len(b) - 1)
for i, x in enumerate(a):
for j, y in enumerate(b):
c[i + j] = c[i + j] + x * y
return c
def _add(a: Sequence[Any], b: Sequence[Any]) -> Poly:
"""Return the coefficientwise sum of two lists of possibly different lengths."""
if len(a) < len(b):
a, b = b, a
return [x + y for x, y in zip(a, b)] + list(a[len(b):])
def karatsuba_mul(a: Sequence[Any], b: Sequence[Any]) -> Poly:
"""Return the full product of two coefficient lists by Karatsuba's method.
Works over every commutative ring. The recursion splits the longer operand
in halves, so its depth is O(log(len(a) + len(b))).
"""
if not a or not b:
return []
if len(a) < len(b):
a, b = b, a
if len(b) <= KARATSUBA_CUTOFF:
return schoolbook_mul(a, b)
zero = a[0] - a[0]
h = len(a) // 2
a0, a1 = list(a[:h]), list(a[h:])
if len(b) <= h:
# Unbalanced: b fits in the low half; multiply each half of a by b.
lo = karatsuba_mul(a0, b)
hi = karatsuba_mul(a1, b)
c = [zero] * (len(a) + len(b) - 1)
for i, v in enumerate(lo):
c[i] = c[i] + v
for i, v in enumerate(hi):
c[i + h] = c[i + h] + v
return c
b0, b1 = list(b[:h]), list(b[h:])
z0 = karatsuba_mul(a0, b0)
z2 = karatsuba_mul(a1, b1)
sa = _add(a0, a1)
sb = _add(b0, b1)
z1 = karatsuba_mul(sa, sb)
for i, v in enumerate(z0):
z1[i] = z1[i] - v
for i, v in enumerate(z2):
z1[i] = z1[i] - v
c = [zero] * (len(a) + len(b) - 1)
for i, v in enumerate(z0):
c[i] = c[i] + v
for i, v in enumerate(z1):
if i + h < len(c):
c[i + h] = c[i + h] + v
for i, v in enumerate(z2):
c[i + 2 * h] = c[i + 2 * h] + v
return c
def _unit_inverse(c: Any) -> Any:
"""Return c^(-1), or raise ValueError if c is not a unit."""
if isinstance(c, bool):
raise TypeError("bool is not a ring element here")
if isinstance(c, int):
if c in (1, -1):
return c
raise ValueError("constant term is not a unit of Z")
if c == 0:
raise ValueError("constant term is zero")
one = c - c + 1
try:
return one / c
except ZeroDivisionError as exc:
raise ValueError("constant term is not a unit") from exc
def precisions(n: int) -> List[int]:
"""Return the precisions 1 = k_s < ... < k_0 = n with k_(i+1) = ceil(k_i / 2)."""
ks = [n]
while ks[-1] > 1:
ks.append((ks[-1] + 1) // 2)
ks.reverse()
return ks
def inverse(f: Sequence[Any], n: int, mul: MulFn = karatsuba_mul) -> Poly:
"""Return [g_0, ..., g_(n-1)] with f g = 1 mod x^n.
Precondition: n >= 0, and f_0 is a unit of the coefficient ring.
Cost: at most 3 M(n) + lower order terms ring operations, where M(n) is the
cost of `mul` on inputs of length n.
"""
if n <= 0:
return []
if not f:
raise ValueError("the zero series has no inverse")
g: Poly = [_unit_inverse(f[0])] # @line L1
m = 1 # @line L2
for m2 in precisions(n)[1:]: # @line L3
d = mul(list(f[:m2]), g)[:m2] # @line L4
while len(d) < m2:
d.append(f[0] - f[0])
t = [-c for c in d[m:m2]] # @line L5
c = mul(g, t)[: m2 - m] # @line L6
while len(c) < m2 - m:
c.append(f[0] - f[0])
g = g + c # @line L7
m = m2
return g # @line L8
def divide_series(a: Sequence[Any], f: Sequence[Any], n: int, mul: MulFn = karatsuba_mul) -> Poly:
"""Return the first n coefficients of a / f, for f_0 a unit."""
if n <= 0:
return []
g = inverse(f, n, mul)
if not a:
return [g[0] - g[0]] * n
q = mul(list(a[:n]), g)[:n]
while len(q) < n:
q.append(g[0] - g[0])
return q
def inverse_naive(f: Sequence[Any], n: int) -> Poly:
"""Return the inverse modulo x^n by the O(n^2) recurrence g_k = -f_0^(-1) sum f_j g_(k-j)."""
if n <= 0:
return []
if not f:
raise ValueError("the zero series has no inverse")
u = _unit_inverse(f[0])
g: Poly = [u]
for k in range(1, n):
s = f[0] - f[0]
for j in range(1, min(k, len(f) - 1) + 1):
s = s + f[j] * g[k - j]
g.append(-(u * s))
return g
if __name__ == "__main__":
print(inverse([1, -1, -1], 8))
test_newton_power_series_inversion.py 197 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/newton-power-series-inversion/. The authors are listed in prov/8P4715.jsonl.
import random
from fractions import Fraction
import pytest
from newton_power_series_inversion import (
divide_series,
inverse,
inverse_naive,
karatsuba_mul,
precisions,
schoolbook_mul,
)
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 oracle on plain ints: truncated schoolbook product.
def trunc_mul(a, b, n, p=None):
c = [0] * n
for i, x in enumerate(a[:n]):
for j, y in enumerate(b[: n - i]):
c[i + j] += x * y
return [v % p for v in c] if p is not None else c
def is_one_mod(c, n):
return c[0] == 1 and all(v == 0 for v in c[1:n])
def test_worked_example_fibonacci():
assert inverse([1, -1, -1], 8) == [1, 1, 2, 3, 5, 8, 13, 21]
assert precisions(8) == [1, 2, 4, 8]
assert precisions(5) == [1, 2, 3, 5]
assert precisions(1) == [1]
def test_geometric_series():
assert inverse([1, -1], 10) == [1] * 10
assert inverse([1, 1], 6) == [1, -1, 1, -1, 1, -1]
def test_edge_cases():
assert inverse([1], 0) == []
assert inverse([5], 0) == []
assert inverse([1], 1) == [1]
assert inverse([-1], 4) == [-1, 0, 0, 0]
assert inverse([Fraction(3)], 3) == [Fraction(1, 3), 0, 0]
# trailing zeros in the input change nothing
assert inverse([1, 2, 0, 0, 0, 0], 5) == inverse([1, 2], 5)
# input longer than the precision: extra terms are ignored
assert inverse([1, 2, 3, 4, 5, 6, 7], 2) == [1, -2]
with pytest.raises(ValueError):
inverse([], 3)
with pytest.raises(ValueError):
inverse([0, 1], 3)
with pytest.raises(ValueError):
inverse([2, 1], 3) # 2 is not a unit of Z
with pytest.raises(ValueError):
inverse([ModP(0, 7), ModP(1, 7)], 3)
def test_random_over_integers():
rng = random.Random(101)
for _ in range(150):
k = rng.randint(1, 40)
f = [rng.choice([1, -1])] + [rng.randint(-50, 50) for _ in range(k - 1)]
n = rng.randint(1, 70)
g = inverse(f, n)
assert g == inverse_naive(f, n)
assert is_one_mod(trunc_mul(f, g, n), n)
@pytest.mark.parametrize("p", [2, 3, 7, 998244353])
def test_random_mod_p(p):
rng = random.Random(p)
for _ in range(60):
k = rng.randint(1, 30)
f0 = rng.randrange(1, p)
raw = [f0] + [rng.randrange(p) for _ in range(k - 1)]
f = [ModP(v, p) for v in raw]
n = rng.randint(1, 80)
g = inverse(f, n)
assert len(g) == n
prod = trunc_mul(raw, [x.v for x in g], n, p)
assert prod[0] == 1 and all(v == 0 for v in prod[1:])
assert [x.v for x in g] == [x.v for x in inverse_naive(f, n)]
def test_random_over_rationals():
rng = random.Random(7)
for _ in range(40):
k = rng.randint(1, 20)
f = [Fraction(rng.randint(1, 9), rng.randint(1, 9)) * rng.choice([1, -1])]
f += [Fraction(rng.randint(-9, 9), rng.randint(1, 9)) for _ in range(k - 1)]
n = rng.randint(1, 40)
g = inverse(f, n)
assert g == inverse_naive(f, n)
assert is_one_mod(trunc_mul(f, g, n), n)
def test_large_precision_uses_karatsuba():
rng = random.Random(3)
p = 998244353
raw = [1] + [rng.randrange(p) for _ in range(599)]
f = [ModP(v, p) for v in raw]
n = 600
g = inverse(f, n)
# check f g = 1 mod x^n at a few random coefficients by direct sums
gv = [x.v for x in g]
for k in [0, 1, 2, 300, 599] + [rng.randrange(n) for _ in range(10)]:
s = sum(raw[j] * gv[k - j] for j in range(k + 1)) % p
assert s == (1 if k == 0 else 0)
def test_karatsuba_against_schoolbook():
rng = random.Random(11)
for _ in range(60):
a = [rng.randint(-99, 99) for _ in range(rng.randint(0, 120))]
b = [rng.randint(-99, 99) for _ in range(rng.randint(0, 120))]
assert karatsuba_mul(a, b) == schoolbook_mul(a, b)
# very different lengths
a = [rng.randint(-9, 9) for _ in range(300)]
b = [rng.randint(-9, 9) for _ in range(40)]
assert karatsuba_mul(a, b) == schoolbook_mul(a, b)
assert karatsuba_mul(b, a) == schoolbook_mul(a, b)
def test_divide_series():
# (1 + x) / (1 - x) = 1 + 2x + 2x^2 + ...
assert divide_series([1, 1], [1, -1], 6) == [1, 2, 2, 2, 2, 2]
assert divide_series([], [1, -1], 3) == [0, 0, 0]
assert divide_series([1, 2, 3], [1], 0) == []
rng = random.Random(5)
for _ in range(40):
a = [rng.randint(-9, 9) for _ in range(rng.randint(1, 20))]
f = [1] + [rng.randint(-9, 9) for _ in range(rng.randint(0, 20))]
n = rng.randint(1, 30)
q = divide_series(a, f, n)
assert trunc_mul(f, q, n) == (a + [0] * n)[:n]
def test_schoolbook_multiplier_gives_same_result():
rng = random.Random(9)
f = [1] + [rng.randint(-5, 5) for _ in range(30)]
assert inverse(f, 50, mul=schoolbook_mul) == inverse(f, 50)
What the next level needs
- the tests did not run in this build