多项式根计算出现微小虚部的问题及解决方法

当使用NumPy计算多项式(x-5)²(x-2)的根时，本该是整数的根5附带了极小的虚部，输出结果如下：

[5.+1.27882372e-07j 5.-1.27882372e-07j 2.+0.00000000e+00j]

重现问题的代码：

import numpy as np
poly=[1,-10,25]  # (x-5)**2
p1 = np.poly1d(poly)
p2 = np.poly1d([1,-2]) # (x-2)
poly2 = p1*p2  # (x-5)**2 * (x-2)
print(poly2.roots) 

print(poly2(5))
print(poly2(2))

解决方案 #

这是数值计算中常见的浮点精度误差导致的，可通过np.real_if_close函数去除这些可忽略的微小虚部。该函数会将虚部与实部的比值小于指定阈值的复数转换为实数。

修改后的代码：

import numpy as np
poly=[1,-10,25]  # (x-5)**2
p1 = np.poly1d(poly)
p2 = np.poly1d([1,-2]) # (x-2)
poly2 = p1*p2  # (x-5)**2 * (x-2)

# 移除微小虚部，tol设置虚部/实部的阈值
formatted_roots = np.real_if_close(poly2.roots, tol=1e-5) 
# 转换为浮点数值格式
roots = formatted_roots.astype(float)
print(roots)

运行后输出为纯实数数组：[5. 5. 2.]

内容的提问来源于stack exchange，提问作者cobp