python - 用一个变量求大量函数的根

Question

我正在使用 Python/numpy/scipy 编写一个小型光线追踪器。表面被建模为二维函数，给出高于法平面的高度。我将寻找射线和表面之间的交点的问题简化为寻找具有一个变量的函数的根。这些函数是连续且连续可微的。

有没有一种方法可以比使用 scipy 根查找器(并且可能使用多个进程)简单地遍历所有函数更有效地做到这一点？

编辑:函数是表示射线的线性函数与表面函数之间的差异，被限制在相交平面上。

Answer 1

以下示例显示了使用二分法并行计算函数 x**(a+1) - b(均具有不同的 a 和 b)的 100 万个副本的根。此处大约需要 12 秒。

import numpy

def F(x, a, b):
    return numpy.power(x, a+1.0) - b

N = 1000000

a = numpy.random.rand(N)
b = numpy.random.rand(N)

x0 = numpy.zeros(N)
x1 = numpy.ones(N) * 1000.0

max_step = 100
for step in range(max_step):
    x_mid = (x0 + x1)/2.0
    F0 = F(x0, a, b)
    F1 = F(x1, a, b)
    F_mid = F(x_mid, a, b)
    x0 = numpy.where( numpy.sign(F_mid) == numpy.sign(F0), x_mid, x0 )
    x1 = numpy.where( numpy.sign(F_mid) == numpy.sign(F1), x_mid, x1 )
    error_max = numpy.amax(numpy.abs(x1 - x0))
    print "step=%d error max=%f" % (step, error_max)
    if error_max < 1e-6: break

基本思想是简单地在变量向量上并行运行求根器的所有常用步骤，使用可以在变量向量和定义个体的等效参数向量上进行评估的函数组件功能。条件语句被替换为掩码和 numpy.where() 的组合。这可以继续下去，直到找到所有根都达到所需的精度，或者直到找到足够多的根，值得将它们从问题中移除，并继续处理排除这些根的较小问题。

我选择解决的函数是任意的，但如果函数表现良好，它会有所帮助；在这种情况下，族中的所有函数都是单调的，并且只有一个正根。此外，对于二分法，我们需要对给出函数不同符号的变量进行猜测，而这些恰好也很容易在这里得出(x0 和 x1 的初始值)。

上面的代码可能使用了最简单的求根器(二分法)，但同样的技术可以很容易地应用于 Newton-Raphson、Ridder 等。求根方法中的条件越少，它就越适合于这个。但是，您将不得不重新实现您想要的任何算法，无法直接使用现有的库根查找器函数。

上面的代码片段是为了清晰而不是速度而编写的。避免某些计算的重复，特别是每次迭代只评估函数一次而不是 3 次，将计算速度加快到 9 秒，如下所示:

...
F0 = F(x0, a, b)
F1 = F(x1, a, b)

max_step = 100
for step in range(max_step):
    x_mid = (x0 + x1)/2.0
    F_mid = F(x_mid, a, b)
    mask0 = numpy.sign(F_mid) == numpy.sign(F0)
    mask1 = numpy.sign(F_mid) == numpy.sign(F1)
    x0 = numpy.where( mask0, x_mid, x0 )
    x1 = numpy.where( mask1, x_mid, x1 )
    F0 = numpy.where( mask0, F_mid, F0 )
    F1 = numpy.where( mask1, F_mid, F1 )
...

相比之下，使用 scipy.bisect() 一次找到一个根大约需要 94 秒:

for i in range(N):
    x_root = scipy.optimize.bisect(lambda x: F(x, a[i], b[i]), x0[i], x1[i], xtol=1e-6)