Partial fraction decomposition in Python
Writes a rational function as a sum of simpler rational terms over factors of its denominator.
- File
partial_fraction_decomposition.py, 106 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/partial-fraction-decomposition/. The authors are listed in prov/P053GE.jsonl.
"""Partial fraction decomposition of a rational function a / b over a field K.
A polynomial is a list of coefficients, lowest degree first, with no trailing
zero; the zero polynomial is []. Coefficients are Fraction, or any field type
with +, -, *, / and == 0.
"""
from typing import Dict, List, Sequence, Tuple
Poly = List[object]
def trim(f: Poly) -> Poly:
f = list(f)
while f and f[-1] == 0:
f.pop()
return f
def p_add(f: Poly, g: Poly) -> Poly:
n = max(len(f), len(g))
return trim([(f[i] if i < len(f) else 0) + (g[i] if i < len(g) else 0) for i in range(n)])
def p_sub(f: Poly, g: Poly) -> Poly:
return p_add(f, [-c for c in g])
def p_mul(f: Poly, g: Poly) -> Poly:
if not f or not g:
return []
h = [f[0] * 0] * (len(f) + len(g) - 1)
for i, a in enumerate(f):
for j, b in enumerate(g):
h[i + j] = h[i + j] + a * b
return trim(h)
def p_divmod(f: Poly, g: Poly) -> Tuple[Poly, Poly]:
"""Quotient and remainder of f by a nonzero g (long division)."""
if not g:
raise ZeroDivisionError("division by the zero polynomial")
r = list(f)
q = [g[-1] * 0] * max(0, len(f) - len(g) + 1)
for k in range(len(f) - len(g), -1, -1):
c = r[k + len(g) - 1] / g[-1]
q[k] = c
for j, b in enumerate(g):
r[k + j] = r[k + j] - c * b
return trim(q), trim(r[: len(g) - 1])
def p_pow(f: Poly, e: int) -> Poly:
out: Poly = [f[-1] / f[-1]] if f else []
for _ in range(e):
out = p_mul(out, f)
return out
def ext_gcd(f: Poly, g: Poly) -> Tuple[Poly, Poly, Poly]:
"""Monic gcd h and s, t with s f + t g = h (extended Euclidean algorithm)."""
r0, r1 = trim(f), trim(g)
one = (r0 or r1)[-1] / (r0 or r1)[-1]
s0, s1, t0, t1 = [one], [], [], [one]
while r1:
q, r = p_divmod(r0, r1)
r0, r1 = r1, r
s0, s1 = s1, p_sub(s0, p_mul(q, s1))
t0, t1 = t1, p_sub(t0, p_mul(q, t1))
lc = r0[-1]
return [c / lc for c in r0], [c / lc for c in s0], [c / lc for c in t0]
def partial_fractions(a: Poly, factors: Sequence[Tuple[Poly, int]]) -> Tuple[Poly, Dict[Tuple[int, int], Poly]]:
"""Decompose a / b with b = prod b_i^e_i given as factors = [(b_i, e_i), ...].
Precondition: each b_i has degree >= 1, each e_i >= 1, and the b_i are
pairwise coprime. Returns the polynomial part q and a dict that maps (i, j)
to a_ij with deg a_ij < deg b_i, such that
a / b = q + sum over i and 1 <= j <= e_i of a_ij / b_i^j.
Cost: O(d^2 + deg(a) d) field operations with classical arithmetic, d = deg b.
"""
b: Poly = [1]
powers = []
for bi, ei in factors:
if len(trim(bi)) < 2 or ei < 1:
raise ValueError("each factor needs degree >= 1 and multiplicity >= 1")
powers.append(p_pow(trim(bi), ei))
one = powers[0][-1] / powers[0][-1] if powers else 1
b = [one]
for pw in powers:
b = p_mul(b, pw) # @line L1
q, a = p_divmod(trim(a), b) # @line L2
parts: Dict[Tuple[int, int], Poly] = {}
for i, (bi, ei) in enumerate(factors): # @line L3
ci, rem = p_divmod(b, powers[i]) # @line L4
g, s, _ = ext_gcd(ci, powers[i]) # @line L5
if len(g) != 1:
raise ValueError("the factors must be pairwise coprime")
r = p_divmod(p_mul(a, s), powers[i])[1] # @line L6
for j in range(ei, 0, -1): # @line L7
r, parts[(i, j)] = p_divmod(r, trim(bi)) # @line L8
return q, parts # @line L9
test_partial_fraction_decomposition.py 202 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/partial-fraction-decomposition/. The authors are listed in prov/P053GE.jsonl.
import random
from fractions import Fraction as Fr
import pytest
from partial_fraction_decomposition import partial_fractions
class GF:
"""Z/pZ for p = 7."""
p = 7
def __init__(self, v):
self.v = v.v if isinstance(v, GF) else v % GF.p
def __add__(self, o):
return GF(self.v + GF(o).v)
__radd__ = __add__
def __sub__(self, o):
return GF(self.v - GF(o).v)
def __rsub__(self, o):
return GF(GF(o).v - self.v)
def __mul__(self, o):
return GF(self.v * GF(o).v)
__rmul__ = __mul__
def __neg__(self):
return GF(-self.v)
def __truediv__(self, o):
o = GF(o)
if o.v == 0:
raise ZeroDivisionError
return GF(self.v * pow(o.v, GF.p - 2, GF.p))
def __eq__(self, o):
return self.v == GF(o).v
def __hash__(self):
return self.v
# independent helpers, written separately from the module under test
def ev(f, x):
s = 0 * x
for c in reversed(f):
s = s * x + c
return s
def mul(f, g):
if not f or not g:
return []
h = [0] * (len(f) + len(g) - 1)
for i, a in enumerate(f):
for j, b in enumerate(g):
h[i + j] = h[i + j] + a * b
while h and h[-1] == 0:
h.pop()
return h
def power(f, e):
out = [1]
for _ in range(e):
out = mul(out, f)
return out
def solve(rows, rhs):
"""Unique solution of a square linear system by Gauss-Jordan elimination."""
n = len(rows)
m = [list(r) + [v] for r, v in zip(rows, rhs)]
for col in range(n):
piv = next(r for r in range(col, n) if m[r][col] != 0)
m[col], m[piv] = m[piv], m[col]
pv = m[col][col]
m[col] = [v / pv for v in m[col]]
for r in range(n):
if r != col and m[r][col] != 0:
fac = m[r][col]
m[r] = [a - fac * b for a, b in zip(m[r], m[col])]
return [m[i][n] for i in range(n)]
def brute(a, factors, field):
"""Solve a = q b + sum a_ij b / b_i^j for the coefficients, as a linear system."""
b = [field(1)]
for bi, ei in factors:
b = mul(b, power(bi, ei))
d = len(b) - 1
dq = max(0, len(a) - 1 - d + 1) if len(a) > d else 0
cols = [] # each column is a polynomial, the coefficient of one unknown
labels = []
for k in range(dq):
cols.append(mul([0] * k + [field(1)], b))
labels.append(("q", k))
for i, (bi, ei) in enumerate(factors):
for j in range(1, ei + 1):
cof = [field(1)]
for i2, (b2, e2) in enumerate(factors):
cof = mul(cof, power(b2, e2 - j if i2 == i else e2))
for k in range(len(bi) - 1):
cols.append(mul([0] * k + [field(1)], cof))
labels.append((i, j, k))
n = len(cols)
size = max(n, len(a))
rows = [[(c[r] if r < len(c) else field(0)) for c in cols] for r in range(size)]
rhs = [(a[r] if r < len(a) else field(0)) for r in range(size)]
sol = solve(rows[:n], rhs[:n]) if n else []
assert all(sum((row[k] * sol[k] for k in range(n)), field(0)) == v for row, v in zip(rows[n:], rhs[n:]))
q, parts = [field(0)] * dq, {}
for lab, v in zip(labels, sol):
if lab[0] == "q":
q[lab[1]] = v
else:
i, j, k = lab
parts.setdefault((i, j), [field(0)] * (len(factors[i][0]) - 1))[k] = v
tr = lambda f: [c for c in f] if not f or f[-1] != 0 else tr(f[:-1])
return tr(q), {key: tr(v) for key, v in parts.items()}
def fill(parts, factors):
return {(i, j): parts.get((i, j), []) for i, (_, e) in enumerate(factors) for j in range(1, e + 1)}
def test_worked_example():
# (x^5 + 2) / ((x - 1)^2 (x^2 + 1))
a = [Fr(2), 0, 0, 0, 0, Fr(1)]
factors = [([Fr(-1), Fr(1)], 2), ([Fr(1), 0, Fr(1)], 1)]
q, parts = partial_fractions(a, factors)
assert q == [Fr(2), Fr(1)]
assert parts[(0, 2)] == [Fr(3, 2)]
assert parts[(0, 1)] == [Fr(1)]
assert parts[(1, 1)] == [Fr(-1, 2), Fr(1)]
def test_edges():
lin = [Fr(-3), Fr(1)]
assert partial_fractions([], [(lin, 2)]) == ([], {(0, 2): [], (0, 1): []})
assert partial_fractions([Fr(5)], [(lin, 1)]) == ([], {(0, 1): [Fr(5)]})
assert partial_fractions([Fr(1), Fr(2)], []) == ([Fr(1), Fr(2)], {})
# numerator a multiple of the denominator: only a polynomial part
q, parts = partial_fractions(mul([Fr(1), Fr(1)], power(lin, 3)), [(lin, 3)])
assert q == [Fr(1), Fr(1)] and all(v == [] for v in parts.values())
with pytest.raises(ValueError):
partial_fractions([Fr(1)], [([Fr(2)], 1)])
with pytest.raises(ValueError):
partial_fractions([Fr(1)], [([Fr(-1), Fr(1)], 1), ([Fr(-2), Fr(1), Fr(1)], 1)]) # x - 1 divides x^2 + x - 2
def random_factors(rng, field, roots):
rs = rng.sample(roots, rng.randint(1, 4))
factors, k = [], 0
while k < len(rs):
g = rng.randint(1, 2)
bi = [field(1)]
for c in rs[k : k + g]:
bi = mul(bi, [field(-c), field(1)])
factors.append((bi, rng.randint(1, 3)))
k += g
if field is Fr and rng.random() < 0.5:
factors.append(([Fr(rng.randint(1, 5)), Fr(0), Fr(1)], rng.randint(1, 2))) # x^2 + c, no rational root
return factors
@pytest.mark.parametrize("field", [Fr, GF])
def test_random_against_linear_algebra_and_points(field):
rng = random.Random(1873 if field is Fr else 77)
roots = list(range(-6, 0)) + list(range(1, 7)) if field is Fr else list(range(7))
for _ in range(500):
factors = random_factors(rng, field, roots)
b = [field(1)]
for bi, ei in factors:
b = mul(b, power(bi, ei))
a = [field(rng.randint(-9, 9)) for _ in range(rng.randint(0, len(b) + 2))]
while a and a[-1] == 0:
a.pop()
q, parts = partial_fractions(a, factors)
for (i, j), aij in parts.items():
assert len(aij) < len(factors[i][0])
bq, bparts = brute(a, factors, field)
assert q == bq
assert fill(parts, factors) == fill(bparts, factors)
pts = range(-12, 13) if field is Fr else range(7)
for x0 in pts:
x = field(x0)
if ev(b, x) == 0:
continue
rhs = ev(q, x)
for (i, j), aij in parts.items():
rhs = rhs + ev(aij, x) / ev(power(factors[i][0], j), x)
assert rhs == ev(a, x) / ev(b, x)
What the next level needs
- the tests did not run in this build