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

我比較舊的機器上的 output 大約是 2.9 秒。 有沒有辦法在 Python 中加速這段代碼？

我的代碼中還有其他地方可以使用一些優化。 當然，我可以創建單獨的帖子，但如果有人也可以發布一些他們認為對優化科學計算的 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 循環，所以事情真的很快。