Quadratic formula in Python
Closed form for the roots of a x^2 + b x + c. Take the larger root without cancellation and the other from the product c/a; with an accurate discriminant both are accurate.
- File
quadratic_formula.py, 65 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/quadratic-formula/. The authors are listed in prov/CQ88BP.jsonl.
"""The quadratic formula in the form that avoids cancellation."""
import math
from dataclasses import dataclass
from fractions import Fraction
from typing import Tuple, Union
Number = Union[float, complex]
@dataclass
class QuadraticResult:
"""The two roots (with multiplicity) and their kind."""
roots: Tuple[Number, ...]
kind: str # "real", "complex", "not quadratic" or "invalid"
def _discriminant(h: float, a: float, c: float) -> float:
"""h*h - a*c computed exactly and rounded once."""
return float(Fraction(h) * Fraction(h) - Fraction(a) * Fraction(c))
def quadratic_roots(a: float, b: float, c: float) -> QuadraticResult:
"""Roots of a x^2 + b x + c with real float coefficients.
Precondition: a, b, c finite floats; for full accuracy every nonzero
coefficient should be at least 2^-1000 times the largest one, so that
the scaling below does not make it subnormal. Real roots are returned
in increasing order; complex roots as a conjugate pair, positive
imaginary part first. Each computed root has relative error below
about 3.5 u (u = 2^-53) with respect to the roots of the stored
coefficients. Cost: a fixed number of operations and one square root.
"""
if not (math.isfinite(a) and math.isfinite(b) and math.isfinite(c)): # @line L1
return QuadraticResult((), "invalid")
if a == 0:
return QuadraticResult((), "not quadratic")
e = math.frexp(max(abs(a), abs(b), abs(c)))[1] # @line L2
a, b, c = math.ldexp(a, -e), math.ldexp(b, -e), math.ldexp(c, -e)
h = b / 2 # @line L3
d = _discriminant(h, a, c)
if d < 0: # @line L4
re = -h / a # @line L5
im = math.sqrt(-d) / a
return QuadraticResult((complex(re, abs(im)), complex(re, -abs(im))), "complex") # @line L6
s = math.sqrt(d) # @line L7
q = -(h + s) if h >= 0 else s - h
if q == 0: # @line L8
return QuadraticResult((0.0, 0.0), "real")
x1, x2 = q / a, c / q # @line L9
return QuadraticResult((min(x1, x2), max(x1, x2)), "real")
def quadratic_roots_naive(a: float, b: float, c: float) -> Tuple[Number, Number]:
"""The textbook formula (-b +- sqrt(b^2 - 4ac)) / (2a), for comparison only."""
d = b * b - 4 * a * c
if d < 0:
s = complex(0.0, math.sqrt(-d))
else:
s = math.sqrt(d)
return ((-b - s) / (2 * a), (-b + s) / (2 * a))
test_quadratic_formula.py 103 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/quadratic-formula/. The authors are listed in prov/CQ88BP.jsonl.
import math
import random
from decimal import Decimal, localcontext
from quadratic_formula import quadratic_roots, quadratic_roots_naive
U = 2.0**-53
def reference(a, b, c):
"""Roots of the stored coefficients by the textbook formula at 120 digits."""
with localcontext() as ctx:
ctx.prec = 120
A, B, C = Decimal(a), Decimal(b), Decimal(c)
d = B * B - 4 * A * C
if d < 0:
re = -B / (2 * A)
im = abs((-d).sqrt() / (2 * A))
return "complex", (re, im)
s = d.sqrt()
r = sorted([(-B - s) / (2 * A), (-B + s) / (2 * A)])
return "real", tuple(r)
def rel(x, y):
return abs(Decimal(x) - y) / abs(y) if y != 0 else abs(Decimal(x))
def test_cancellation_example():
r = quadratic_roots(1.0, -1e8, 1.0)
assert r.kind == "real" and r.roots == (1e-08, 1e8)
naive = quadratic_roots_naive(1.0, -1e8, 1.0)
assert naive[0] == 7.450580596923828e-09 # 25 percent wrong
def test_close_roots_exact_discriminant():
b, c = -(2 + 2.0**-26), 1 + 2.0**-26
r = quadratic_roots(1.0, b, c)
assert r.roots == (1.0, 1 + 2.0**-26)
naive = quadratic_roots_naive(1.0, b, c)
assert naive[0] == naive[1] == 1 + 2.0**-27 # the discriminant rounded to 0
def test_complex_pair():
r = quadratic_roots(1.0, 1.0, 1.0)
assert r.kind == "complex"
assert r.roots == (complex(-0.5, 0.8660254037844386), complex(-0.5, -0.8660254037844386))
r = quadratic_roots(-2.0, 0.0, -8.0) # pure imaginary 2i, -2i with a < 0
assert r.roots == (2j, -2j)
def test_double_root_and_zero_root():
assert quadratic_roots(1.0, -2.0, 1.0).roots == (1.0, 1.0)
assert quadratic_roots(3.0, 0.0, 0.0).roots == (0.0, 0.0)
assert quadratic_roots(2.0, -6.0, 0.0).roots == (0.0, 3.0)
def test_not_quadratic_and_invalid():
assert quadratic_roots(0.0, 1.0, 1.0).kind == "not quadratic"
assert quadratic_roots(math.nan, 1.0, 1.0).kind == "invalid"
assert quadratic_roots(1.0, math.inf, 1.0).kind == "invalid"
def test_no_overflow_with_huge_coefficients():
big = math.ldexp(1.0, 1000)
r = quadratic_roots(big, -3 * big, 2 * big)
assert r.roots == (1.0, 2.0)
naive = quadratic_roots_naive(big, -3 * big, 2 * big)
assert math.isnan(naive[0]) or math.isinf(naive[0])
tiny = math.ldexp(1.0, -1000)
r = quadratic_roots(tiny, -3 * tiny, 2 * tiny)
assert r.roots == (1.0, 2.0)
assert quadratic_roots_naive(tiny, -3 * tiny, 2 * tiny) != (1.0, 2.0)
def test_vieta():
rng = random.Random(7)
for _ in range(200):
a, b, c = (rng.uniform(-10, 10) for _ in range(3))
r = quadratic_roots(a, b, c)
x, y = r.roots
assert abs(x + y + b / a) <= 1e-12 * (abs(x) + abs(y) + abs(b / a))
assert abs(x * y - c / a) <= 1e-12 * (abs(x * y) + abs(c / a))
def test_random_against_decimal_reference():
rng = random.Random(20261107)
for _ in range(600):
a, b, c = (rng.choice((-1, 1)) * rng.uniform(1, 2) * 10.0 ** rng.uniform(-8, 8) for _ in range(3))
r = quadratic_roots(a, b, c)
kind, ref = reference(a, b, c)
assert r.kind == kind
if kind == "real":
for got, want in zip(r.roots, ref):
assert rel(got, want) <= Decimal(4 * U)
else:
z = r.roots[0]
assert rel(z.real, ref[0]) <= Decimal(4 * U)
assert rel(z.imag, ref[1]) <= Decimal(4 * U)
assert r.roots[1] == z.conjugate()
What the next level needs
- the tests did not run in this build