在不使用循环的情况下计算 numpy 中多个多项式的根

Question

I can use the polyfit() method with a 2D array as input, to calculate polynomials on multiple data sets in a fast manner. After getting these multiple polynomials, I want to calculate the roots of all of these polynomials, in a fast manner.

There is numpy.roots() method for finding the roots of a single polynomial but this method does not work with 2D inputs (meaning multiple polynomials). I am working with millions of polynomials, so I would like to avoid looping over all polynomials using a for loop, map or comprehension because it takes minutes in that case. I would prefer a vectoral numpy operation or series of vectoral operations.

An example code for inefficient calculation:

POLYNOMIAL_COUNT = 1000000

# Create a polynomial of second order with coefficients 2, 3 and 4
coefficients = np.array([[2,3,4]])

# Let's say we have the same polynomial multiple times, represented as a 2D array.
# In reality the polynomial coefficients will be different from each other,
# but they will be the same order.
coefficients = coefficients.repeat(POLYNOMIAL_COUNT, axis=0)

# Calculate roots of these same-order polynomials.
# Looping here takes too much time.
roots = []
for i in range(POLYNOMIAL_COUNT):
    roots.append(np.roots(coefficients[i]))

Is there a way to find the roots of multiple same-order polynomials using numpy, but without looping?

Answer 1

For the special case of polynomials up to the fourth order, you can solve in a vectorized manner. Anything higher than that does not have an analytical solution, so requires iterative optimization, which is fundamentally unlikely to be vectorizable since different rows may require a different number of iterations. As @John Coleman suggests , you might be able to get away with using the same number of steps for each one, but will likely have to sacrifice accuracy to do so.

That being said, here is an example of how to vectorize the second order case:

d = coefficients[:, 1:-1]**2 - 4.0 * coefficients[:, ::2].prod(axis=1, keepdims=True)
roots = -0.5 * (coefficients[:, 1:-1] + [1, -1] * np.emath.sqrt(d)) / coefficients[:, :1]

If I got the order of the coefficients wrong, replace coefficients[:, :1] with coefficients[:, -1:] in the denominator of the last assignment. Using np.emath.sqrt is nice because it will return a complex128 result automatically when your discriminant d is negative anywhere, and normal float64 result for all real roots.

You can implement a third order solution or a fourth order solution in a similar manner.