Source code for foxes.utils.cubic_roots

import numpy as np



[docs]
def cubic_roots(a0, a1, a2, a3=None):
    """
    Calculate real roots of polynomials of degree 3.

    Convention:
    f(x) = a[3]*x**3 + a[2]*x**2 + a[1]*x + a[0]

    In contrast to numpy's "root" function
    this works fast for an array of polynomials,
    so you spare yourself looping over them.

    Source: https://github.com/opencv/opencv/blob/master/modules/calib3d/src/polynom_solver.cpp

    Parameters
    ----------
    a0: numpy.ndarray
        The coefficients a[0]
    a1: numpy.ndarray
        The coefficients a[1]
    a2: numpy.ndarray
        The coefficients a[2]
    a3: numpy.ndarray
        The coefficients a[3], or None for ones

    Returns
    -------
    roots: numpy.ndarray
        The real roots of the polynomial,
        shape: (n_a0, 3). If one root only
        the two last columns will be np.nan

    :group: utils

    """

    N = len(a0)
    out = np.full([N, 3], np.nan)

    # Calculate the normalized form x^3 + a2 * x^2 + a1 * x + a0 = 0
    b_a = a2 if a3 is None else a2 / a3
    b_a2 = b_a * b_a
    c_a = a1 if a3 is None else a1 / a3
    d_a = a0 if a3 is None else a0 / a3

    # Solve the cubic equation
    Q = (3 * c_a - b_a2) / 9
    R = (9 * b_a * c_a - 27 * d_a - 2 * b_a * b_a2) / 54
    Q3 = Q * Q * Q
    D = Q3 + R * R
    b_a_3 = (1.0 / 3.0) * b_a

    sel = Q == 0.0
    if np.any(sel):
        o = out[sel, 0]

        sel2 = R == 0.0
        if np.any(sel2):
            o[sel2] = -b_a_3[sel][sel2]

        if np.any(~sel2):
            o[~sel2] = np.pow(2 * R[sel][~sel2], 1 / 3.0) - b_a_3[sel][~sel2]

        out[sel, 0] = o

    sel = D <= 0.0
    if np.any(sel):
        # Three real roots
        theta = np.arccos(R[sel] / np.sqrt(-Q3[sel]))
        sqrt_Q = np.sqrt(-Q[sel])

        out[sel, 0] = 2 * sqrt_Q * np.cos(theta / 3.0) - b_a_3[sel]
        out[sel, 1] = 2 * sqrt_Q * np.cos((theta + 2 * np.pi) / 3.0) - b_a_3[sel]
        out[sel, 2] = 2 * sqrt_Q * np.cos((theta + 4 * np.pi) / 3.0) - b_a_3[sel]

    return out



def test_cubic_roots(roots, a0, a1, a2, a3=None, tol=1.0e-12):
    """
    Test the cubic roots results

    Parameters
    ----------
    roots: numpy.ndarray
        The roots to test, shape: (n_a0, 3)
    a0: numpy.ndarray
        The coefficients a[0]
    a1: numpy.ndarray
        The coefficients a[1]
    a2: numpy.ndarray
        The coefficients a[2]
    a3: numpy.ndarray
        The coefficients a[3], or None for ones

    """

    N = len(a0)
    for n in range(N):
        c0 = a0[n]
        c1 = a1[n]
        c2 = a2[n]
        c3 = a3[n]

        print(f"Polynomial {n}: a = {(c0,c1,c2,c3)}")

        rts = np.unique(roots[n])
        rts = rts[~np.isnan(rts)]

        for x in rts:
            f = c0 + c1 * x + c2 * x**2 + c3 * x**3
            ok = np.abs(f) <= tol

            print(f"  root x = {x}: f(x) = {f}     {'OK' if ok else 'FAILED'}")

            if not ok:
                raise Exception("NOT OK!")

        if len(rts) == 0:
            print("  no real roots.")


if __name__ == "__main__":
    N = 100
    a0 = np.random.uniform(-10.0, 10.0, N)
    a1 = np.random.uniform(-10.0, 10.0, N)
    a2 = np.random.uniform(-10.0, 10.0, N)
    a3 = np.random.uniform(1.0, 10.0, N)

    roots = cubic_roots(a0, a1, a2, a3)

    test_cubic_roots(roots, a0, a1, a2, a3)