Algorithm Search⌘K

Cardano's formula in Python

The three roots of a cubic in closed form. Shift to remove the square term, solve a quadratic for u^3, take one cube root u, and form the roots u - p/(3u) with the three choices of u.

File
cardano_formula.py, 48 lines
Level
I0 not checked yet
Toolchain
Python 3.14, pytest 9
Test framework
pytest
Run the tests
python3 -m pytest
cardano_formula.pyRawGitHub
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/cardano-formula/. The authors are listed in prov/XGZ50G.jsonl.

"""Cardano's formula for the three roots of a cubic, in complex arithmetic."""

import cmath
from typing import Tuple

OMEGA = complex(-0.5, 3 ** 0.5 / 2)  # exp(2 pi i / 3)


def cardano(a: complex, b: complex, c: complex, d: complex) -> Tuple[complex, complex, complex]:
    """Return the three roots, with multiplicity, of a x^3 + b x^2 + c x + d.

    Precondition: a != 0. The coefficients may be real or complex. The
    partner of the cube root u is taken as -p / (3u), and the root of the
    resolvent quadratic of larger modulus is used, so no cancellation to 0
    occurs there. Cost: a fixed number of complex operations, one square
    root and one cube root, O(1) space.
    """
    if a == 0:
        raise ValueError("a must be nonzero")
    s, C, D = b / (3 * a), c / a, d / a  # @line L1
    p, q = C - 3 * s * s, D - s * C + 2 * s ** 3  # @line L2
    r = cmath.sqrt(q * q / 4 + p ** 3 / 27)  # @line L3
    w = -q / 2 + r  # @line L4
    if abs(-q / 2 - r) > abs(w):
        w = -q / 2 - r
    if w == 0:  # @line L5
        return (-s, -s, -s)
    u = complex(w) ** (1 / 3)  # @line L6
    roots = []
    for k in range(3):  # @line L7
        uk = u * OMEGA ** k  # @line L8
        tk = uk - p / (3 * uk)  # @line L9
        roots.append(tk - s)  # @line L10
    return tuple(roots)  # @line L11


def polish(a: complex, b: complex, c: complex, d: complex, x: complex, steps: int = 2) -> complex:
    """A few Newton steps on the cubic from x, to remove rounding errors of the formula."""
    for _ in range(steps):
        fx = ((a * x + b) * x + c) * x + d
        dfx = (3 * a * x + 2 * b) * x + c
        if dfx == 0:
            break
        x = x - fx / dfx
    return x
test_cardano_formula.py 99 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/cardano-formula/. The authors are listed in prov/XGZ50G.jsonl.

import cmath
import random

import pytest

from cardano_formula import cardano, polish


def match(got, want, tol):
    """Match two multisets of complex numbers greedily within tol."""
    rest = list(want)
    for z in got:
        j = min(range(len(rest)), key=lambda i: abs(rest[i] - z))
        if abs(rest[j] - z) > tol:
            return False
        rest.pop(j)
    return True


def test_cardano_example():
    roots = cardano(1, 0, 6, -20)
    assert abs(roots[0] - 2) < 1e-14
    assert match(roots, [2, -1 + 3j, -1 - 3j], 1e-13)


def test_three_real_roots_need_complex_intermediates():
    roots = cardano(1, -6, 11, -6)
    assert match(roots, [1, 2, 3], 1e-13)
    assert [round(z.real) for z in roots] == [3, 1, 2]
    assert all(abs(z.imag) < 1e-15 for z in roots)


def test_double_and_triple_roots():
    assert match(cardano(1, 0, -3, 2), [1, 1, -2], 1e-7)
    assert cardano(1, 0, 0, 0) == (0, 0, 0)
    assert match(cardano(1, -3, 3, -1), [1, 1, 1], 1e-5)  # (x - 1)^3: about u^(1/3)


def test_sign_choice_avoids_false_triple_root():
    roots = cardano(1, 0, 1e-8, 1)
    want = [-1, cmath.exp(1j * cmath.pi / 3), cmath.exp(-1j * cmath.pi / 3)]
    assert match(roots, want, 1e-8)


def test_independent_cube_roots_pitfall():
    r = cmath.sqrt(100 + 8)
    u, v = (10 + r) ** (1 / 3), complex(10 - r) ** (1 / 3)
    t = u + v
    assert abs(t ** 3 + 6 * t - 20) > 1  # not a root
    assert abs(cardano(1, 0, 6, -20)[0] - 2) < 1e-14


def test_zero_coefficients_and_complex_input():
    assert match(cardano(2, 0, 0, -16), [2, 2 * cmath.exp(2j * cmath.pi / 3),
                                         2 * cmath.exp(-2j * cmath.pi / 3)], 1e-13)
    want = [1j, 2, -1 - 1j]
    a = 3
    b = -a * sum(want)
    c = a * (want[0] * want[1] + want[0] * want[2] + want[1] * want[2])
    d = -a * want[0] * want[1] * want[2]
    assert match(cardano(a, b, c, d), want, 1e-12)


def test_a_zero_is_rejected():
    with pytest.raises(ValueError):
        cardano(0, 1, 2, 3)


def test_roots_of_different_sizes_and_polishing():
    r = [1e-6, 1.0, 1000.0]
    b, c, d = -sum(r), r[0] * r[1] + r[0] * r[2] + r[1] * r[2], -r[0] * r[1] * r[2]
    small = min(cardano(1, b, c, d), key=abs)
    assert 1e-6 < abs(small - 1e-6) / 1e-6 < 1e-4  # the shift loses digits
    assert abs(polish(1, b, c, d, small) - 1e-6) / 1e-6 < 1e-12


def test_random_cubics_with_known_roots():
    rng = random.Random(31)
    for _ in range(600):
        if rng.random() < 0.5:
            want = [complex(rng.uniform(-10, 10), 0) for _ in range(3)]
        else:
            x = rng.uniform(-10, 10)
            z = complex(rng.uniform(-10, 10), rng.uniform(0.5, 10))
            want = [complex(x, 0), z, z.conjugate()]
        gap = min(abs(want[i] - want[j]) for i in range(3) for j in range(i))
        if gap < 0.1:
            continue
        a = rng.choice([1, -2.5, 7])
        b = -a * sum(want)
        c = a * (want[0] * want[1] + want[0] * want[2] + want[1] * want[2])
        d = -a * want[0] * want[1] * want[2]
        got = cardano(a, b.real, c.real, d.real)
        assert match(got, want, 1e-9 * 10 / gap)
        for z in got:
            assert abs(((a * z + b) * z + c) * z + d) < 1e-9 * abs(a) * 1e3

What the next level needs

  • the tests did not run in this build

Read the full page about Cardano's formula