Newton's method for square root in Python
Approximates the square root of a by repeatedly averaging a guess x with a/x. Converges from every positive start, quadratically, with an exact error formula.
- File
newton_square_root.py, 59 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-square-root/. The authors are listed in prov/MBGG7T.jsonl.
"""Newton's method for the square root (Heron's method), with range reduction."""
import math
from dataclasses import dataclass
from typing import Any, List
@dataclass
class SqrtResult:
"""The computed square root, why the loop ended, and the work done."""
root: float
reason: str # "exact input", "converged", "kmax reached" or "invalid"
iterations: int
divisions: int
def newton_sqrt(a: float, max_iter: int = 60) -> SqrtResult:
"""Square root of a float by Newton's method x <- (x + m/x)/2.
Precondition: a is a float. Negative numbers and NaN give "invalid";
0, -0.0 and +inf are returned unchanged. The argument is reduced
exactly to a = m * 4^e with 1/2 <= m < 2, the iteration runs on m from
x = (1 + m)/2 and stops as soon as an iterate fails to decrease. The
result is within one unit in the last place of the true root. Cost:
one division per pass, at most 6 passes for binary64.
"""
if math.isnan(a) or a < 0: # @line L1
return SqrtResult(math.nan, "invalid", 0, 0)
if a == 0 or math.isinf(a):
return SqrtResult(a, "exact input", 0, 0)
f, big_e = math.frexp(a) # @line L2
if big_e % 2 == 0:
m, e = f, big_e // 2
else:
m, e = 2 * f, (big_e - 1) // 2
x = (1 + m) / 2 # @line L3
for k in range(1, max_iter + 1): # @line L4
y = (x + m / x) / 2 # @line L5
if y >= x: # @line L6
return SqrtResult(math.ldexp(x, e), "converged", k, k)
x = y # @line L7
return SqrtResult(math.ldexp(x, e), "kmax reached", max_iter, max_iter) # @line L8
def heron_iterates(a: Any, x0: Any, n: int) -> List[Any]:
"""The first n Newton iterates x <- (x + a/x)/2 from x0 > 0, in any field type.
Works with Fraction (exact) and Decimal (current context). Cost: n divisions.
"""
xs = [x0]
x = x0
for _ in range(n):
x = (x + a / x) / 2
xs.append(x)
return xs
test_newton_square_root.py 87 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/newton-square-root/. The authors are listed in prov/MBGG7T.jsonl.
import math
import random
from decimal import Decimal, localcontext
from fractions import Fraction
from newton_square_root import heron_iterates, newton_sqrt
def test_worked_example():
xs = heron_iterates(2.0, 1.0, 5)
assert xs == [1.0, 1.5, 1.4166666666666665, 1.4142156862745097, 1.4142135623746899, 1.414213562373095]
r = newton_sqrt(2.0)
assert r.reason == "converged" and r.iterations == 5 and r.divisions == 5
assert r.root == 1.414213562373095 # one unit below the correctly rounded value
assert math.sqrt(2.0) - r.root == math.ulp(1.4)
def test_fraction_iterates_and_identity():
xs = heron_iterates(Fraction(2), Fraction(1), 4)
assert xs == [1, Fraction(3, 2), Fraction(17, 12), Fraction(577, 408), Fraction(665857, 470832)]
# t_k = (x_k - 3)/(x_k + 3) satisfies t_{k+1} = t_k^2 exactly for a = 9
xs = heron_iterates(Fraction(9), Fraction(1), 6)
ts = [(x - 3) / (x + 3) for x in xs]
assert ts[0] == Fraction(-1, 2)
for k in range(6):
assert ts[k + 1] == ts[k] ** 2
def test_monotone_after_first_step():
for x0 in (1e-3, 0.5, 1.0, 7.0, 1e3):
xs = heron_iterates(Fraction(5), Fraction(x0), 6)
for k in range(1, 6):
assert xs[k] ** 2 >= 5 and xs[k + 1] <= xs[k]
def test_observed_order_two():
with localcontext() as ctx:
ctx.prec = 300
a = Decimal(3)
r = a.sqrt()
xs = heron_iterates(a, Decimal(1), 8)
e = [abs(x - r) for x in xs]
for k in range(3, 8):
q = (e[k + 1].ln() - e[k].ln()) / (e[k].ln() - e[k - 1].ln())
assert abs(q - 2) < 0.05
# the constant: e_{k+1}/e_k^2 tends to 1/(2 sqrt a)
assert abs(e[8] / e[7] ** 2 - 1 / (2 * r)) < Decimal("1e-10")
def test_edge_inputs():
assert newton_sqrt(-1.0).reason == "invalid"
assert math.isnan(newton_sqrt(math.nan).root)
assert newton_sqrt(math.inf).root == math.inf
assert newton_sqrt(0.0).root == 0.0
z = newton_sqrt(-0.0).root
assert z == 0.0 and math.copysign(1.0, z) == -1.0
assert newton_sqrt(5e-324).root == math.sqrt(5e-324)
big = newton_sqrt(1.7976931348623157e308)
assert abs(big.root - math.sqrt(1.7976931348623157e308)) <= math.ulp(big.root)
def test_kmax_reached():
r = newton_sqrt(3.0, max_iter=1)
assert r.reason == "kmax reached" and r.iterations == 1 and r.divisions == 1
assert r.root == 1.7321428571428572
def test_perfect_squares_exact():
rng = random.Random(11)
for _ in range(2000):
k = rng.randint(1, 2**26)
assert newton_sqrt(float(k * k)).root == k
def test_random_within_one_ulp():
rng = random.Random(20261107)
for _ in range(5000):
a = math.ldexp(rng.uniform(0.5, 1.0), rng.randint(-1070, 1020))
r = newton_sqrt(a)
assert r.reason == "converged" and r.iterations <= 6
with localcontext() as ctx:
ctx.prec = 40
exact = Decimal(a).sqrt()
assert abs(Decimal(r.root) - exact) <= Decimal(math.ulp(r.root))
What the next level needs
- the tests did not run in this build