Python：有没有比使用 numpy 的 polyfromroots 更快的替代方法？

Question

使用Python，我循环遍历大量根集（本质上表示为二维 numpy 数组），并每次形成具有根集的一元多项式（称为根）作为根，因此：

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

然而，当根集的数量很大（大约 40000）时，我的代码运行运行缓慢，即使每组根本身很小（仅包含 4 个根）。

举个例子，如果我运行

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))

我相对较旧的机器上的输出大约是 2.9 秒。有没有办法在 Python 内加速这段代码？

我的代码中还有其他地方可以使用一些优化。当然，我可以创建单独的 posts，但如果有人也可以 post 一些他们认为对优化 Python 科学计算代码有用的资源，这将对我（和其他人）有所帮助。

Answer 1

Sympy 可以预先计算给定根数的通用公式：

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])

调用help(np_poly_4_roots)显示其源代码：

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])

这已经运行得更快了，但是没有完全矢量化。但是您可以使用此源手动创建矢量化版本：

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)])

对于具有所有 4 个根的输入，这可以以矢量化方式执行：

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)

因为现在没有 Python for 循环了，事情真的很快。

Python：有没有比使用 numpy 的 polyfromroots 更快的替代方法？

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

python

optimization

numpy

scientific-computing