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

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

我可以使用 polyfit() 方法以二维数组作为输入，以快速方式计算多个数据集上的多项式。得到这些多项式后，我想快速计算所有这些多项式的根。

有 numpy.roots() 方法可用于查找单个多项式的根，但此方法不适用于 2D 输入（即多个多项式）。我正在处理数百万个多项式，因此我想避免使用 for 循环、映射或推导式遍历所有多项式，因为在这种情况下需要几分钟。我更喜欢向量 numpy 操作或一系列向量操作。

低效计算的示例代码：

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

有没有办法使用 numpy 找到多个同阶多项式的根，但不循环？

解决方法

对于四阶以下的多项式的特殊情况，您可以以向量化的方式求解。高于此值的任何值都没有解析解，因此需要迭代优化，这从根本上不太可能可向量化，因为不同的行可能需要不同的迭代次数。作为 @John Coleman suggests，您可能可以避免对每个步骤使用相同数量的步骤，但可能必须为此牺牲准确性。

话虽如此，以下是如何矢量化二阶情况的示例：

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]

如果我弄错了系数的顺序，请将上次分配的分母中的 coefficients[:,:1] 替换为 coefficients[:,-1:]。使用 np.emath.sqrt 很好，因为当您的判别式 complex128 在任何地方为负时，它会自动返回 d 结果，而所有实根的正常 float64 结果。

您可以以类似的方式实现 third order solution 或 fourth order solution。