根据条件在numpy数组的元素中进行数学运算的有效方法

Question

我正在尝试优化我的python代码.当我尝试根据每个元素值将函数应用于numpy数组时,出现了一个瓶颈.例如,我有一个包含数千个元素的数组,我为一个大于公差的值应用一个函数,为其余的函数应用另一个函数(泰勒系列).我做掩蔽但仍然很慢,至少我调用了6400万次以下的功能.

EPSILONZETA = 1.0e-6
ZETA1_12 = 1.0/12.0
ZETA1_720 = 1.0/720.0

def masked_condition_zero(array, tolerance):
    """ Return the indices where values are lesser (and greater) than tolerance
    """
    # search indices where array values < tolerance
    indzeros_ = np.where(np.abs(array) < tolerance)[0]

    # create mask
    mask_ = np.ones(np.shape(array), dtype=bool)

    mask_[[indzeros_]] = False

    return (~mask_, mask_) 

def bernoulli_function1(zeta):
    """ Returns the Bernoulli function of zeta, vector version
    """
    # get the indices according to condition
    zeros_, others_ = masked_condition_zero(zeta, EPSILONZETA)

    # create an array filled with zeros
    fb_ = np.zeros(np.shape(zeta))

    # Apply the original function to the values greater than EPSILONZETA
    fb_[others_] = zeta[others_]/(np.exp(zeta[others_])-1.0)  

    # computes series for zeta < eps
    zeta0_ = zeta[zeros_]
    zeta2_ = zeta0_ *  zeta0_
    zeta4_ =  zeta2_ * zeta2_
    fb_[zeros_] = 1.0 - 0.5*zeta0_ + ZETA1_12 * zeta2_ - ZETA1_720 * zeta4_
    return fb_

现在假设你有一个带有负浮动和正浮动的数组zeta,它在每个循环中发生变化,延伸到2 ^ 26次迭代,你想每次计算fbernoulli_function1(zeta).

还有更好的解决方案吗？

Answer 1

问题的基本结构是：

\n\n

def foo(zeta):\n    result = np.empty_like(zeta)\n    I = condition(zeta)\n    nI = ~I\n    result[I] = func1(zeta[I])\n    result[nI] = func2(zeta[nI])\n

\n\n

看起来你的多项式表达式根本可以被计算，但这是“例外”，当太接近 0zeta时的回退计算。zeta

\n\n

如果两个函数都可以针对 a 求值zeta，则可以使用 where：

\n\n

np.where(condition(zeta), func1(zeta), func2(zeta))\n

\n\n

这是以下内容的简化版本：

\n\n

def foo(zeta):\n    result = np.empty_like(zeta)\n    I = condition(zeta)\n    nI = ~I\n    v1 = func1(zeta)\n    v2 = func2(zeta)\n    result[I] = v1[I]\n    result[nI] = v2[nI]\n

\n\n

另一种选择是将一个函数应用于所有值，而另一个函数仅应用于“例外”。

\n\n

def foo(zeta):\n    result = func2(zeta)\n    I = condition(zeta)\n    result[I] = func1[zeta[I]]\n

\n\n

当然，相反 - result = func1(zeta); result[nI]=func2[zeta]。

\n\n

在我的简短测试中，func1大约func2需要相同的时间。

\n\n

masked_condition_zero也需要时间，但是更简单np.abs(array) < tolerance（而且它是~J）可以将时间减少一半。

\n\n

让我们比较一下分配策略

\n\n

def foo(zeta, J, nJ):\n    result = np.empty_like(zeta)\n    result[J] = fun1(zeta[J])\n    result[nJ] = fun2(zeta[nJ])\n    return result\n

\n\n

zeta[J]对于为 10% 的样本zeta，一些样本时间为：

\n\n

In [127]: timeit foo(zeta, J, nJ)\n10000 loops, best of 3: 55.7 \xc2\xb5s per loop\n\nIn [128]: timeit result=fun2(zeta); result[J]=fun1(zeta[J])\n10000 loops, best of 3: 49.2 \xc2\xb5s per loop\n\nIn [129]: timeit np.where(J, fun1(zeta),fun2(zeta))\n10000 loops, best of 3: 73.4 \xc2\xb5s per loop\n\nIn [130]: timeit result=fun1(zeta); result[nJ]=fun2(zeta[nJ])\n10000 loops, best of 3: 60.7 \xc2\xb5s per loop\n

\n\n

第二种情况是最快的，因为运行fun1较少的值可以补偿索引的额外成本zeta[J]。索引成本和功能评估成本之间需要权衡。像这样的布尔索引比切片更昂贵。如果采用其他价值组合，时间安排可能会走向另一个方向。

\n\n

这看起来像是一个可以逐渐减少时间的问题，但我没有看到任何突破策略可以将时间缩短一个数量级。

\n