Python: is there a faster alternative to using numpy's polyfromroots?-CodePudding

Using Python, I am looping over a large collection of sets of roots (represented as a 2-dimensional numpy array, essentially), and forming each time the monic polynomial having the set of roots (called roots) as roots, as such:

poly = np.polynomial.polynomial.polyfromroots(roots)

However, my code runs slowly when the number of sets of roots is large (about 40000), even though each set of roots is itself small (containing only 4 roots).

As an example, if I run

import timeit

SETUP_CODE = '''
import numpy as np'''

TEST_CODE = '''
N, n = 40000, 4
collection_roots = np.random.random((N,n))   1.j*np.random.random((N,n))
polynomials = np.zeros((N, n   1), dtype=complex)
for k in range(N):
    roots = collection_roots[k, :]
    polynomials[k, :] = np.polynomial.polynomial.polyfromroots(roots)'''

print(timeit.timeit(setup=SETUP_CODE,
                    stmt=TEST_CODE,
                    number=1))

the output on my relatively old machine is about 2.9 seconds. Is there a way to speed up this piece of code, within Python?

There are other places in my code that could use some optimization. Of course, I could create separate posts, but it would help me (and others) if someone could also post some resources they found useful for optimizing Python code for scientific computing.

CodePudding user response：

Sympy can precalculate the general formula for a given number of roots:

from sympy import Symbol, Product, poly, lambdify

num_roots = 4
x = Symbol('x')
roots = [Symbol(f'r{i}') for i in range(num_roots)]
prod = 1
for ri in roots:
    prod *= (x - ri)
print(prod)
print(poly(prod, x).all_coeffs()[::-1])
np_poly_4_roots = lambdify(roots, poly(prod, x).all_coeffs()[::-1])

np_poly_4_roots(*[1, 2, 3, 4])

Calling help(np_poly_4_roots) shows its source code:

def _lambdifygenerated(r0, r1, r2, r3):
    return ([r0*r1*r2*r3,
             -r0*r1*r2 - r0*r1*r3 - r0*r2*r3 - r1*r2*r3,
             r0*r1   r0*r2   r0*r3   r1*r2   r1*r3   r2*r3,
             -r0 - r1 - r2 - r3,
             1])

This already works a bit faster, but isn't fully vectorized. But you could use this source to manually create a vectorized version:

def fast_poly_4_roots(r):
    r0, r1, r2, r3 = r[:, 0], r[:, 1], r[:, 2], r[:, 3]
    return np.array([r0 * r1 * r2 * r3,
                     -r0 * r1 * r2 - r0 * r1 * r3 - r0 * r2 * r3 - r1 * r2 * r3,
                     r0 * r1   r0 * r2   r0 * r3   r1 * r2   r1 * r3   r2 * r3,
                     -r0 - r1 - r2 - r3,
                     np.ones_like(r0)])

For an input with all 4 roots, this can be executed in a vectorized way:

N, n = 4000, 4
collection_roots = np.random.random((N,n))   1.j*np.random.random((N,n))

polynomials = fast_poly_4_roots(collection_roots)

As now there aren't Python for-loops anymore, things are really fast.