pooltool.ptmath.roots.quartic

Overview

Classes

QuarticSolver

Enum where members are also (and must be) strings

Function

solve_quartics(ps, solver)

Returns the smallest positive and real root for each quartic polynomial.

solve_many_numerical(p)

Solve multiple polynomial equations using companion matrix eigenvalues

solve_many(ps)

Solve multiple quartic equations using analytical solutions when possible

solve(a, b, c, d, e)

-

evaluate(p, val)

-

instability(p)

Range is from [0, inf], 0 is most stable

numeric(p)

-

analytic(p)

Calculate a quartic’s roots using the closed-form solution

Classes

class pooltool.ptmath.roots.quartic.QuarticSolver(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)[source]

Bases: pooltool.utils.strenum.StrEnum

Functions

pooltool.ptmath.roots.quartic.solve_quartics(ps: numpy.typing.NDArray[numpy.float64], solver: QuarticSolver = QuarticSolver.HYBRID) numpy.typing.NDArray[numpy.float64][source]

Returns the smallest positive and real root for each quartic polynomial.

Parameters:
  • ps (numpy.typing.NDArray[numpy.float64]) -- A mx5 array of polynomial coefficients, where m is the number of equations. The columns are in the order a, b, c, d, e, where these coefficients make up the quartic polynomial equation at^4 + bt^3 + ct^2 + dt + e = 0.

  • solver (QuarticSolver) -- The method used to calculate the roots. See pooltool.ptmath.roots.quartic.QuarticSolver.

Returns:

An array of shape m. Each value is the smallest root that is real and positive. If no such root exists (e.g. all roots have complex), then np.inf is returned.

Return type:

roots

pooltool.ptmath.roots.quartic.solve_many_numerical(p: numpy.typing.NDArray[numpy.float64]) numpy.typing.NDArray[numpy.complex128][source]

Solve multiple polynomial equations using companion matrix eigenvalues

This is a vectorized implementation of numpy.roots that can solve multiple polynomials in a vectorized fashion. The solution is taken from this wonderful stackoverflow answer: https://stackoverflow.com/a/35853977

Parameters:

p (numpy.typing.NDArray[numpy.float64]) -- A mxn array of polynomial coefficients, where m is the number of equations and n-1 is the order of the polynomial. If n is 5 (4th order polynomial), the columns are in the order a, b, c, d, e, where these coefficients make up the polynomial equation at^4 + bt^3 + ct^2 + dt + e = 0

Return type:

numpy.typing.NDArray[numpy.complex128]

Notes

  • Not yet amenable to numbaization (0.56.4). Problem is the numba implementation of np.linalg.eigvals, which only supports 2D arrays, but the strategy here is to pass np.lingalg.eigvals as a vectorized 3D array. Nevertheless, here is a numba implementation that is just slightly slower (7% slower) than this function:

    n = p.shape[-1] A = np.zeros(p.shape[:1] + (n - 1, n - 1), dtype=np.complex128) A[:, 1:, :-1] = np.eye(n - 2) p0 = np.copy(p[:, 0]).reshape((-1, 1)) A[:, 0, :] = -p[:, 1:] / p0 roots = np.zeros((p.shape[0], n - 1), dtype=np.complex128) for i in range(p.shape[0]):

    roots[i, :] = np.linalg.eigvals(A[i, :, :])

    return roots

pooltool.ptmath.roots.quartic.solve_many(ps: numpy.typing.NDArray[numpy.float64]) numpy.typing.NDArray[numpy.complex128][source]

Solve multiple quartic equations using analytical solutions when possible

Closed-form analytical solutions exist for the quartic polynomial equation, but can suffer from severe numerical instability. Fortunately, the quality of an analytically calculated roots can be determined by plugging them back into the quartic and ensuring they evaluate the function to 0.

This function calculates roots to a quartic by analytically solving the quartic polynomials. If the roots are inaccurate, an isomorphic polynomial is solved analytically. If those roots are also inaccurate, the roots are solved using the companion matrix eigenvalue method, which is very reliable, but slower.

Parameters:

p -- A mx5 array of polynomial coefficients, where m is the number of equations. The columns are in the order a, b, c, d, e, where these coefficients make up the polynomial equation at^4 + bt^3 + ct^2 + dt + e = 0

Return type:

numpy.typing.NDArray[numpy.complex128]

pooltool.ptmath.roots.quartic.solve(a: float, b: float, c: float, d: float, e: float) numpy.typing.NDArray[numpy.complex128][source]
Return type:

numpy.typing.NDArray[numpy.complex128]

pooltool.ptmath.roots.quartic.evaluate(p: numpy.typing.NDArray[numpy.complex128], val: complex) complex[source]
Return type:

complex

pooltool.ptmath.roots.quartic.instability(p: numpy.typing.NDArray[numpy.complex128]) float[source]

Range is from [0, inf], 0 is most stable

Return type:

float

pooltool.ptmath.roots.quartic.numeric(p: numpy.typing.NDArray[numpy.complex128]) numpy.typing.NDArray[numpy.complex128][source]
Return type:

numpy.typing.NDArray[numpy.complex128]

pooltool.ptmath.roots.quartic.analytic(p: numpy.typing.NDArray[numpy.complex128]) numpy.typing.NDArray[numpy.complex128][source]

Calculate a quartic’s roots using the closed-form solution

This function was created with the help of sympy.

To start, I solved the general quartic polynomial roots:

>>> from sympy import symbols, Eq, solve
>>> x, a, b, c, d, e = symbols('x a b c d e')
>>> general_solution = solve(a*x**4 + b*x**3 + c*x**2 + d*x + e, x)

This yields 4 expressions, one for each root. Each expression is piecewise conditional, where if the following equality is true, the first expression is used, and otherwise the second expression is used.

>>> general_solution[0].args[0][1]
Eq(e/a - b*d/(4*a**2) + c**2/(12*a**2), 0)

So in total there are 8 expressions, 2 for each root, and the expression used for each root is determined based on whether the above equality holds true. These expressions are huge, so to better digest them, I used the following common subexpression elimination:

>>> from sympy import cse
>>> cse(
>>>     [
>>>         general_solution[0].args[0][0],
>>>         general_solution[0].args[1][0],
>>>         general_solution[1].args[0][0],
>>>         general_solution[1].args[1][0],
>>>         general_solution[2].args[0][0],
>>>         general_solution[2].args[1][0],
>>>         general_solution[3].args[0][0],
>>>         general_solution[3].args[1][0],
>>>     ]
>>> )

Then I used a vim macro to convert these subexpressions into the lines of python code you see below.

Return type:

numpy.typing.NDArray[numpy.complex128]