PSLQ algorithm in Python
Iterative algorithm that finds integer relations among real numbers.
- File
pslq_algorithm.py, 104 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/pslq-algorithm/. The authors are listed in prov/GH5331.jsonl.
from fractions import Fraction
from math import floor
from typing import Optional, Sequence, Union
Number = Union[int, Fraction]
def _nearest(v: Fraction) -> int:
"""The integer nearest to v; a half is rounded up."""
return floor(v + Fraction(1, 2))
def pslq_counted(
x: Sequence[Number], gamma2: Fraction = Fraction(2), max_iterations: int = 100000
) -> tuple[list[int], int, list[Fraction]]:
"""PSLQ in exact rational arithmetic.
Input: a vector x of n >= 2 rationals, not all zero, and gamma2 = gamma^2 > 4/3.
Output: (m, iterations, bounds), where m is a nonzero integer vector with
sum m_i x_i = 0, iterations is the number of swaps, and bounds[t] is
max_j ||p_j*||^2 before iteration t; every integer relation r of x has
||r||^2 >= 1 / bounds[t] for every t.
The matrix H of the usual description is represented by the Gram-Schmidt data
of the rows p_i of A P_x: H[j][j]^2 = big_b[j] and H[i][j] / H[j][j] = mu[i][j].
Raises ValueError for bad input and RuntimeError if max_iterations is exceeded.
"""
n = len(x)
if n < 2:
raise ValueError("need at least two numbers")
xs = [Fraction(v) for v in x]
norm2 = sum((v * v for v in xs), Fraction(0))
if norm2 == 0:
raise ValueError("x must not be the zero vector")
if gamma2 <= Fraction(4, 3):
raise ValueError("gamma^2 must exceed 4/3")
a = [[int(i == j) for j in range(n)] for i in range(n)] # @line L1
bmat = [[int(i == j) for j in range(n)] for i in range(n)] # B = A^-1, relations are columns
y = list(xs) # y = x B
for j in range(n):
if y[j] == 0:
return [bmat[i][j] for i in range(n)], 0, []
# @line L2: Gram-Schmidt of p_i = e_i - (x_i / <x, x>) x, from the Gram matrix of the p_i
gram = [[Fraction(int(i == j)) - xs[i] * xs[j] / norm2 for j in range(n)] for i in range(n)]
mu = [[Fraction(0)] * n for _ in range(n)]
big_b = [Fraction(0)] * n # big_b[n-1] is 0 throughout
for i in range(n):
for j in range(min(i, n - 1)):
s = gram[i][j] - sum((mu[j][k] * mu[i][k] * big_b[k] for k in range(j)), Fraction(0))
mu[i][j] = s / big_b[j]
if i < n - 1:
big_b[i] = gram[i][i] - sum((mu[i][k] ** 2 * big_b[k] for k in range(i)), Fraction(0))
def hermite_reduce() -> None:
for i in range(1, n): # @line L10
for j in range(i - 1, -1, -1):
t = _nearest(mu[i][j]) # @line L11
if t == 0:
continue
a[i] = [p - t * q for p, q in zip(a[i], a[j])]
for k in range(j):
mu[i][k] -= t * mu[j][k]
mu[i][j] -= t
y[j] += t * y[i] # @line L12
for row in bmat:
row[j] += t * row[i]
hermite_reduce() # @line L3
bounds: list[Fraction] = []
iterations = 0
while True: # @line L4
for j in range(n): # @line L5
if y[j] == 0:
return [bmat[i][j] for i in range(n)], iterations, bounds
if iterations >= max_iterations:
raise RuntimeError("no relation found within max_iterations")
bounds.append(max(big_b[: n - 1]))
r = max(range(n - 1), key=lambda j: (gamma2 ** (j + 1) * big_b[j], -j)) # @line L6
k = r + 1
a[r], a[k] = a[k], a[r] # @line L7
y[r], y[k] = y[k], y[r]
for row in bmat:
row[r], row[k] = row[k], row[r]
# @line L8: the swap formulas of LLL, with big_b[n-1] = 0
u = mu[k][r]
new_r = big_b[k] + u * u * big_b[r]
mu[k][r] = u * big_b[r] / new_r
big_b[k] = big_b[r] * big_b[k] / new_r
big_b[r] = new_r
for j in range(r):
mu[r][j], mu[k][j] = mu[k][j], mu[r][j]
for i in range(k + 1, n):
t = mu[i][k]
mu[i][k] = mu[i][r] - u * t
mu[i][r] = t + mu[k][r] * mu[i][k]
hermite_reduce() # @line L9
iterations += 1
def pslq(x: Sequence[Number], gamma2: Fraction = Fraction(2)) -> list[int]:
"""A nonzero integer vector m with sum m_i x_i = 0, found by PSLQ."""
return pslq_counted(x, gamma2)[0]
test_pslq_algorithm.py 139 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/pslq-algorithm/. The authors are listed in prov/GH5331.jsonl.
import itertools
import random
from fractions import Fraction
from math import gcd, isqrt
import pytest
from pslq_algorithm import pslq, pslq_counted
def dot(u, v):
return sum((a * b for a, b in zip(u, v)), 0)
def shortest_relation_sq(x, bound_sq):
"""min ||m||^2 over nonzero integer m with <m, x> = 0 and ||m||^2 <= bound_sq.
Every such m has |m_i| <= isqrt(bound_sq), so the box search is exhaustive.
"""
c = isqrt(bound_sq)
best = None
for m in itertools.product(range(-c, c + 1), repeat=len(x)):
if any(m) and dot(m, x) == 0:
s = dot(m, m)
if s <= bound_sq and (best is None or s < best):
best = s
return best
def integer_scale(x):
"""The primitive integer vector on the line of x (same relations)."""
x = [Fraction(v) for v in x]
den = 1
for v in x:
den = den * v.denominator // gcd(den, v.denominator)
z = [int(v * den) for v in x]
g = 0
for v in z:
g = gcd(g, v)
return [v // g for v in z]
def check(x, gamma2=Fraction(2), brute_limit=40000):
m, iterations, bounds = pslq_counted(x, gamma2)
assert any(m) and dot(m, x) == 0
s = dot(m, m)
# max_j ||p_j*||^2 never increases, and starts at most 1
assert all(b2 <= b1 for b1, b2 in zip(bounds, bounds[1:]))
assert all(b <= 1 for b in bounds)
# the iteration bound: tau^(2 I) <= ||x||^(2 (n - 1)) for the primitive integer x
tau2 = 4 * gamma2 / (4 + gamma2)
z = integer_scale(x)
assert tau2 ** iterations <= dot(z, z) ** (len(x) - 1)
if (2 * isqrt(s) + 1) ** len(x) <= brute_limit:
best = shortest_relation_sq(x, s)
assert best is not None and best <= s
# the lower bound of every iteration: ||r||^2 >= 1 / max_j ||p_j*||^2
assert all(best * b >= 1 for b in bounds)
return m, iterations
def test_worked_example():
m, iterations, bounds = pslq_counted([3, 5, 7])
assert m == [-1, 2, -1] and iterations == 3
assert bounds == [Fraction(74, 83), Fraction(74, 83), Fraction(22, 83)]
check([3, 5, 7])
def test_not_always_shortest():
x = [Fraction(-5, 2), -6, -7]
m = pslq(x)
assert m == [0, 7, -6]
assert shortest_relation_sq(x, dot(m, m)) == 9
assert dot([2, -2, 1], x) == 0
def test_two_numbers():
m = pslq([355, -113])
assert m in ([113, 355], [-113, -355])
rng = random.Random(1)
for _ in range(100):
p, q = rng.randint(-500, 500), rng.randint(1, 500)
if p == 0:
continue
m = pslq([Fraction(p, q), 1])
g = gcd(p, q)
assert m in ([q // g, -p // g], [-q // g, p // g])
def test_zero_entries():
assert pslq([0, 4, 5]) == [1, 0, 0]
assert pslq([2, 3, 0]) == [0, 0, 1]
assert pslq_counted([1, 1])[0] in ([1, -1], [-1, 1])
def test_known_relation():
# x = (1, 1/2, 1/3, 1/4): many small relations, PSLQ must return a short one
check([1, Fraction(1, 2), Fraction(1, 3), Fraction(1, 4)])
check([12, 18, 30, 42, 66])
def test_hard_rational_approximation():
# close to (1, sqrt 2, sqrt 3): only long relations exist
x = [1, Fraction(14142136, 10**7), Fraction(17320508, 10**7)]
m, iterations, bounds = pslq_counted(x)
assert all(dot(m, m) * b >= 1 for b in bounds)
check(x)
def test_random_vectors():
rng = random.Random(20251106)
for n in range(2, 7):
for _ in range(40 if n <= 4 else 12):
hi = 30 if n <= 4 else 6
x = [Fraction(rng.randint(-hi, hi), rng.randint(1, 4)) for _ in range(n)]
if not any(x):
continue
check(x)
def test_other_gamma():
rng = random.Random(9)
for gamma2 in (Fraction(3, 2), Fraction(4)):
for _ in range(30):
x = [rng.randint(-40, 40) for _ in range(rng.randint(2, 4))]
if any(x):
check(x, gamma2)
def test_errors():
with pytest.raises(ValueError):
pslq([5])
with pytest.raises(ValueError):
pslq([0, 0, 0])
with pytest.raises(ValueError):
pslq([1, 2], Fraction(4, 3))
What the next level needs
- the tests did not run in this build