"""
rootfinding.py
This code is a collection of functions and classes designed to find the roots of polynomial
equations. In mathematics, finding the roots of a polynomial means determining the values of x
that make the polynomial equal to zero. This is a common problem in various fields like engineering,
physics, and computer science.
The main purpose of this code is to implement the Bairstow's method, which is an algorithm for
finding complex roots of polynomials. It takes as input the coefficients of a polynomial and
initial guesses for the roots, and outputs the calculated roots of the polynomial.
The code starts by importing necessary modules and defining some utility classes and functions.
The main algorithm is implemented in the pbairstow_even function. This function takes three inputs:
1. A list of coefficients representing the polynomial
2. A list of initial guesses for the roots
3. An optional Options object to control the algorithm's behavior
The output of pbairstow_even is a tuple containing:
1. A list of the calculated roots
2. The number of iterations performed
3. A boolean indicating whether the algorithm successfully converged
The algorithm works by iteratively refining the initial guesses for the roots. It uses a technique
called "suppression" to improve the accuracy of each root estimate. The process continues until
either the desired accuracy is achieved or the maximum number of iterations is reached.
Some important parts of the code include:
- The initial_guess function, which generates starting points for the algorithm
- The suppress function, which helps improve the accuracy of root estimates
- The horner function, which efficiently evaluates polynomials
- The delta function, which calculates adjustments to the root estimates
The code also includes helper functions for polynomial evaluation and manipulation, such as
horner_eval and horner_backward.
Overall, this code provides a sophisticated tool for solving polynomial equations, even when the
roots are complex numbers. It's designed to be efficient and accurate, making it useful for
applications that require finding roots of high-degree polynomials.
"""
from functools import reduce
from itertools import accumulate
from math import sqrt
from typing import Any, List, Sequence, Tuple, Union
from mywheel.robin import Robin
from .matrix2 import Matrix2
from .vector2 import Vector2
Num = Union[float, complex]
# The class "Options" defines default values for maximum iterations, tolerance, and individual
# tolerance.
[docs]
class Options:
"""
Configuration options for root-finding algorithms.
This class provides control parameters for iterative root-finding methods:
- max_iters: Maximum number of iterations allowed (default: 2000)
- tolerance: Convergence tolerance for the algorithm (default: 1e-12)
- tol_ind: Individual root tolerance threshold (default: 1e-15)
These parameters allow fine-tuning of the algorithm's behavior and stopping criteria.
"""
max_iters: int = 2000
tolerance: float = 1e-12
tol_ind: float = 1e-15
# tol_suppress: float = 1e-1
[docs]
def delta(vA: Vector2, vr: Vector2, vp: Vector2) -> Vector2:
"""Calculate adjustment vector for Bairstow's method.
The `delta` function computes the correction vector used in Bairstow's method to update
root estimates. It solves a 2x2 linear system derived from polynomial division to find
the optimal adjustment to current root estimates.
:param vA: Residual vector (A,B) from polynomial division
:type vA: Vector2
:param vr: Current root estimate vector (r,q)
:type vr: Vector2
:param vp: Vector used in suppression calculations (p,s)
:type vp: Vector2
:return: Correction vector to adjust root estimates
:rtype: Vector2
.. svgbob::
┌ ┐ -1 ┌ ┐
│ r⋅p + s p │ │A│
│ │ ⋅ │ │
│ q⋅p s │ │B│
└ ┘ └ ┘
Examples:
>>> d = delta(Vector2(1, 2), Vector2(-2, 0), Vector2(4, 5))
>>> print(d)
<0.2, 0.4>
"""
r_coeff, q_coeff = vr.x, vr.y
p_coeff, s_coeff = vp.x, vp.y
mp = Matrix2(
Vector2(s_coeff, -p_coeff),
Vector2(-p_coeff * q_coeff, p_coeff * r_coeff + s_coeff),
)
return mp.mdot(vA) / mp.det() # 6 mul's + 2 div's
[docs]
def suppress_old(vA: Vector2, vA1: Vector2, vri: Vector2, vrj: Vector2) -> None:
"""Original implementation of zero suppression in Bairstow's method.
This function modifies the residual vectors vA and vA1 to suppress the influence
of other roots (vrj) when estimating the current root (vri). This helps prevent
interference between root estimates during iteration.
Note: This is the original implementation that modifies vectors in-place.
The newer version returns modified vectors instead.
:param vA: Current residual vector (A,B)
:type vA: Vector2
:param vA1: First derivative residual vector (A1,B1)
:type vA1: Vector2
:param vri: Current root estimate being refined (ri,qi)
:type vri: Vector2
:param vrj: Another root estimate that might interfere (rj,qj)
:type vrj: Vector2
Reference:
D. C. Handscomb, Computation of the latent roots of a Hessenberg matrix
by Bairsow's method, Computer Journal, 5 (1962), pp. 139-141.
Examples:
>>> vA = Vector2(3, 3)
>>> vA1 = Vector2(1, 2)
>>> vri = Vector2(-2, 0)
>>> vrj = Vector2(4, 5)
>>> suppress_old(vA, vA1, vri, vrj)
>>> dr = delta(vA, vri, vA1)
>>> print(dr)
<-16.78082191780822, 1.4383561643835616>
"""
A_val, B_val = vA.x, vA.y
A1_val, B1_val = vA1.x, vA1.y
vp = vri - vrj
r_coeff, q_coeff = vri.x, vri.y
p_coeff, s_coeff = vp.x, vp.y
f_val = r_coeff * p_coeff + s_coeff
qp_val = q_coeff * p_coeff
e_val = f_val * s_coeff - qp_val * p_coeff
a_val = A_val * s_coeff - B_val * p_coeff
b_val = B_val * f_val - A_val * qp_val
c_val = A1_val * e_val - a_val
d_val = B1_val * e_val - b_val - a_val * p_coeff
vA._x = a_val * e_val
vA._y = b_val * e_val
vA1._x = c_val * s_coeff - d_val * p_coeff
vA1._y = d_val * f_val - c_val * qp_val
# return delta(vA, vri, Vector2(vA1._x, -vA1._y))
[docs]
def suppress(
vA: Vector2, vA1: Vector2, vri: Vector2, vrj: Vector2
) -> Tuple[Vector2, Vector2]:
"""Improved zero suppression for Bairstow's method.
This function calculates modified residual vectors that account for interference
from other roots in the system. It uses matrix operations to efficiently compute
the suppression terms, providing better numerical stability than the original version.
:param vA: Current residual vector (A,B)
:type vA: Vector2
:param vA1: First derivative residual vector (A1,B1)
:type vA1: Vector2
:param vri: Current root estimate being refined (ri,qi)
:type vri: Vector2
:param vrj: Another root estimate that might interfere (rj,qj)
:type vrj: Vector2
:return: Tuple of modified residual vectors (vA, vA1)
:rtype: Tuple[Vector2, Vector2]
Reference:
D. C. Handscomb, Computation of the latent roots of a Hessenberg matrix
by Bairsow's method, Computer Journal, 5 (1962), pp. 139-141.
Examples:
>>> vA = Vector2(3, 3)
>>> vA1 = Vector2(1, 2)
>>> vri = Vector2(-2, 0)
>>> vrj = Vector2(4, 5)
>>> vA, vA1 = suppress(vA, vA1, vri, vrj)
>>> dr = delta(vA, vri, vA1)
>>> print(dr)
<-16.78082191780822, 1.4383561643835616>
"""
vp = vri - vrj
r_coeff, q_coeff = vri.x, vri.y
p_coeff, s_coeff = vp.x, vp.y
m_adjoint = Matrix2(
Vector2(s_coeff, -p_coeff),
Vector2(-p_coeff * q_coeff, p_coeff * r_coeff + s_coeff),
)
e_val = m_adjoint.det()
va = m_adjoint.mdot(vA)
vc = vA1 * e_val - va
vc._y -= va._x * p_coeff
va *= e_val
va1 = m_adjoint.mdot(vc)
return va, va1
[docs]
def horner_eval_f(coeffs: Sequence[Num], zval: Num) -> Num:
"""Evaluate polynomial using Horner's method (functional version).
This function computes the value of a polynomial at a given point using
Horner's method, which is more efficient than direct evaluation. It uses
Python's reduce function for a concise implementation.
:param coeffs: List of polynomial coefficients in descending order of degree
:type coeffs: List
:param zval: Point at which to evaluate the polynomial (can be complex)
:return: Value of the polynomial at zval
:rtype: Same type as zval (float or complex)
Examples:
>>> coeffs = [1, -8, -72, 382, 727, -2310]
>>> horner_eval_f(coeffs, 3)
960
"""
return reduce(lambda res, coeff: res * zval + coeff, coeffs)
# n n - 1
# P(z) = c ⋅ z + c ⋅ z + ... + c
# 0 1 n
#
# P(z) = P (z) ⋅ ⎛z - z ⎞ + A
# 1 ⎝ val⎠
[docs]
def horner_eval(coeffs: Sequence[Num], zval: Num) -> Tuple[Any, List[Num]]:
"""Evaluate polynomial and return intermediate coefficients.
This function uses Horner's method to evaluate a polynomial and also returns
the intermediate coefficients that result from the synthetic division process.
These coefficients can be used for further computations like derivatives.
:param coeffs: List of polynomial coefficients in descending order of degree
:type coeffs: List
:param zval: Point at which to evaluate the polynomial (can be complex)
:return: Tuple containing:
- Value of polynomial at zval
- List of intermediate coefficients from synthetic division
:rtype: Tuple[Any, List]
Examples:
>>> coeffs = [1, -8, -72, 382, 727, -2310]
>>> horner_eval(coeffs, 3)
(960, [1, -5, -87, 121, 1090, 960])
>>> horner_eval(coeffs, 3+0j)
((960+0j), [1, (-5+0j), (-87+0j), (121+0j), (1090+0j), (960+0j)])
"""
coeffs = list(accumulate(coeffs, lambda res, coeff: res * zval + coeff))
return coeffs[-1], coeffs
#
# ⎛ 2 ⎞
# P(x) = P (x) ⋅ ⎝x - r ⋅ x - q⎠ + A ⋅ x + B
# 1
#
# Note: P(x) becomes the quotient after calling this function
[docs]
def horner(coeffs: List[float], degree: int, vr: Vector2) -> Vector2:
"""Evaluate quadratic polynomial factor and return remainder.
This specialized version of Horner's method evaluates a polynomial divided by
a quadratic factor (x² - r·x - q). It returns the linear remainder (A·x + B)
and modifies the coefficients array to contain the quotient polynomial.
:param coeffs: List of polynomial coefficients in descending order
:type coeffs: List[float]
:param degree: Degree of the polynomial (must be ≥ 2)
:type degree: int
:param vr: Vector representing quadratic factor coefficients (r,q)
:type vr: Vector2
:return: Remainder vector (A,B)
:rtype: Vector2
Examples:
>>> coeffs = [1, -8, -72, 382, 727, -2310]
>>> vp = horner(coeffs, 5, Vector2(-1, 6)) # x^2 + x - 6
>>> coeffs
[1, -9, -57, 385, 0, 0]
>>> coeffs = [1, -8, -72, 382, 727, -2310]
>>> vp = horner(coeffs, 5, Vector2(2, 3)) # x^2 - 2x - 3
>>> coeffs
[1, -6, -81, 202, 888, -1704]
"""
for i in range(0, degree - 1):
coeffs[i + 1] += coeffs[i] * vr.x
coeffs[i + 2] += coeffs[i] * vr.y
return Vector2(coeffs[degree - 1], coeffs[degree])
[docs]
def initial_guess(coeffs: List[float]) -> List[Vector2]:
"""Generate initial root estimates for Bairstow's method.
This function creates reasonable starting points for the root-finding algorithm
by distributing estimates around a circle in the complex plane. The circle's
center and radius are determined from the polynomial's coefficients.
:param coeffs: List of polynomial coefficients in descending order
:type coeffs: List[float]
:return: List of initial root estimates as Vector2 objects (r,q pairs)
:rtype: List[Vector2]
Examples:
>>> h = [10.0, 34.0, 75.0, 94.0, 150.0, 94.0, 75.0, 34.0, 10.0]
>>> vr0s = initial_guess(h)
"""
degree: int = len(coeffs) - 1
center: float = -coeffs[1] / (degree * coeffs[0])
poly_c: Num = horner_eval_f(coeffs, center)
radius: float = pow(abs(poly_c), 1.0 / degree)
quad_term: float = center * center + radius * radius
degree //= 2
degree *= 2 # make even
from .aberth import COS_PI_VDC2_TABLE
temp = iter(radius * COS_PI_VDC2_TABLE[i + 1] for i in range(0, degree // 2))
return [Vector2(2 * (center + t), -(quad_term + 2 * center * t)) for t in temp]
[docs]
def pbairstow_even(
coeffs: List[float], vrs: List[Vector2], options: Options = Options()
) -> Tuple[List[Vector2], int, bool]:
r"""Parallel implementation of Bairstow's root-finding method.
This function implements a parallel version of Bairstow's method for finding
all roots of a polynomial simultaneously. It works by iteratively improving
estimates of quadratic factors of the polynomial.
:param coeffs: List of polynomial coefficients in descending order
:type coeffs: List[float]
:param vrs: Initial estimates for quadratic factors (as Vector2 pairs)
:type vrs: List[Vector2]
:param options: Configuration parameters for the algorithm
:type options: Options
:return: Tuple containing:
- Final root estimates
- Number of iterations performed
- Convergence status (True if converged)
:rtype: Tuple[List[Vector2], int, bool]
.. svgbob::
┌ ┐(new) ┌ ┐ ┌-------------------┐ -1 ┌ ┐
│rᵢ│ │rᵢ│ │ A'₁⋅rᵢ + B'₁ A'₁│ │A│
│ │ │ │ - │ │ ⋅ │ │
│qᵢ│ │qᵢ│ │ A'₁⋅qᵢ B'₁│ │B│
└ ┘ └ ┘ └-------------------┘ └ ┘
where
┌ ┐ ┌ ┐ m ┌--------------------┐ -1 ┌ ┐
│A'₁│ │A₁│ _____ │ pᵢⱼ⋅rᵢ + sᵢⱼ pᵢⱼ │ │A│
│ │ = │ │ - ╲ │ │ ⋅ │ │
│B'₁│ │B₁│ ╱ │ pᵢⱼ⋅qᵢ sᵢⱼ │ │B│
└ ┘ └ ┘ ‾‾‾‾‾ └--------------------┘ └ ┘
j ≠ i
and
┌ ┐ ┌ ┐ ┌ ┐
│pᵢⱼ│ │rᵢ│ │rⱼ│
│ │ = │ │ - │ │
│sᵢⱼ│ │qᵢ│ │qⱼ│
└ ┘ └ ┘ └ ┘
Examples:
>>> h = [10.0, 34.0, 75.0, 94.0, 150.0, 94.0, 75.0, 34.0, 10.0]
>>> vr0s = initial_guess(h)
>>> vrs, niter, found = pbairstow_even(h, vr0s)
>>> print(found)
True
"""
num_factors = len(vrs)
degree = len(coeffs) - 1
converged = [False] * num_factors
robin = Robin(num_factors)
for niter in range(options.max_iters):
tolerance = 0.0
for i, (vri, ci) in enumerate(zip(vrs, converged)):
if ci:
continue
coeffs1 = coeffs.copy()
vA = horner(coeffs1, degree, vri)
tol_i = max(abs(vA.x), abs(vA.y))
if tol_i < options.tol_ind:
converged[i] = True
continue
vA1 = horner(coeffs1, degree - 2, vri)
tolerance = max(tol_i, tolerance)
for j in robin.exclude(i):
vA, vA1 = suppress(vA, vA1, vri, vrs[j])
vrs[i] -= delta(vA, vri, vA1)
if tolerance < options.tolerance:
return vrs, niter, True
return vrs, options.max_iters, False
[docs]
def roots_from_quadratic(vr: Vector2) -> Tuple[complex, complex]:
"""
Extract the two roots of a quadratic factor x^2 - r*x - q = 0.
:param vr: Vector2 with x=r, y=q representing x^2 - r*x - q
:return: The two roots (may be real or complex conjugate pair)
Examples:
>>> r1, r2 = roots_from_quadratic(Vector2(0, 1))
>>> abs(r1 - 1j) < 1e-14
True
>>> abs(r2 + 1j) < 1e-14
True
"""
r = vr.x
q = vr.y
disc = r * r + 4.0 * q
if disc >= 0.0:
sqrt_disc = sqrt(disc)
return ((r + sqrt_disc) / 2.0, (r - sqrt_disc) / 2.0)
sqrt_disc = sqrt(-disc)
return (complex(r / 2.0, sqrt_disc / 2.0), complex(r / 2.0, -sqrt_disc / 2.0))
[docs]
def poly_from_quadratic_factors(vrs: List[Vector2]) -> List[float]:
"""
Reconstruct a monic polynomial from its quadratic factors.
Each factor x^2 - r*x - q is converted to its roots, then all roots are
reconstructed with Leja ordering via poly_from_roots.
:param vrs: Quadratic factors from Bairstow's method
:return: Monic polynomial coefficients (highest degree first)
Examples:
>>> vrs = [Vector2(0.0, 1.0)] # x^2 - 1
>>> coeffs = poly_from_quadratic_factors(vrs)
>>> coeffs
[1.0, 0.0, -1.0]
"""
if not vrs:
return [1.0]
from .aberth import poly_from_roots
all_roots: List[complex] = []
for vr in vrs:
r1, r2 = roots_from_quadratic(vr)
all_roots.append(r1)
all_roots.append(r2)
return poly_from_roots(all_roots)
[docs]
def find_rootq(vr: Vector2) -> Tuple[Num, Num]:
"""Solve quadratic equation x² - r·x - q = 0.
This function finds the roots of a quadratic equation represented by
the Vector2 vr, where vr.x is r and vr.y is q in the equation above.
It handles both real and complex roots appropriately.
:param vr: Vector containing quadratic coefficients (r,q)
:type vr: Vector2
:return: Tuple containing the two roots (real or complex)
:rtype: Tuple[Num, Num]
Examples:
>>> vr = find_rootq(Vector2(5, -6))
>>> print(vr)
(3.0, 2.0)
"""
half_radius = vr.x / 2
discriminant = half_radius * half_radius + vr.y
if discriminant < 0:
root1 = half_radius + sqrt(-discriminant) * 1j
else:
root1 = half_radius + (
sqrt(discriminant) if half_radius >= 0 else -sqrt(discriminant)
)
root2 = -vr.y / root1
return root1, root2