根据条件对numpy数组的元素进行计算的有效方法

Question

我正在尝试优化我的python代码。当我试图根据每个元素值将一个函数应用到numpy数组时，就会出现一个瓶颈。例如，我有一个包含上千个元素的数组，我为大于公差的值应用了一个函数，为其余的元素应用了另一个函数(Taylor级数)。我做了掩蔽，但仍然缓慢，至少我调用了6千4百万次以下函数。

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

问题的基本结构是：

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

看起来您的多项式表达式可以在所有zeta上进行计算，但它是“异常”，当zeta太接近于0时，它是回退计算。

如果可以为zeta计算这两个函数，则可以在以下位置使用：

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

这是简化后的版本：

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

另一个选项是将一个函数应用于所有值，另一个只应用于“异常”。

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

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

在我短暂的时间测试中，func1和func2几乎同时进行。

masked_condition_zero也需要时间，但是简单的np.abs(array) < tolerance (以及它的~J)会将其削减一半。

让我们比较一下分配策略。

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

对于zeta[J]为完整zeta的10%的示例，一些示例时间如下：

In [127]: timeit foo(zeta, J, nJ)
10000 loops, best of 3: 55.7 µs per loop

In [128]: timeit result=fun2(zeta); result[J]=fun1(zeta[J])
10000 loops, best of 3: 49.2 µs per loop

In [129]: timeit np.where(J, fun1(zeta),fun2(zeta))
10000 loops, best of 3: 73.4 µs per loop

In [130]: timeit result=fun1(zeta); result[nJ]=fun2(zeta[nJ])
10000 loops, best of 3: 60.7 µs per loop

第二种情况是最快的，因为在较少的值上运行fun1会补偿索引zeta[J]的额外成本。在索引的成本和功能评估的成本之间有一种权衡。像这样的布尔索引比切片更昂贵。与其他的混合值，时间可以走相反的方向。

这看起来像一个问题，你可以削减时间，但我没有看到任何断断续续的挽歌，会削减一个数量级的时间。

Answer 2

where命令在索引到数组中的速度比较慢。这可能会更快。

fb_ = np.zeros_like(zeta)
nonZero= zeta > ZETA_TOLERANCE
zero = ~nonZero
fb_[zero] = function1(zeta[zero])
fb_[nonZero] = function2(zeta[nonZero])

编辑：我意识到我的原始版本是制作两个相同数组的副本。这个新版本应该更快一些。

Answer 3

您可以使用numba [1](http://numba.pydata.org/) )，它是一个用于处理numpy的jit编译器。

from numba import jit
@jit
def bernoulli_function_fill(zeta, fb_):
    for i in xrange(len(zeta)):
        if np.abs(zeta[i])>EPSILONZETA:
            fb_[i] = zeta[i]/(np.exp(zeta[i])-1.0)
        else:
            zeta0_ = zeta[i]
            zeta2_ = zeta0_ *  zeta0_
            zeta4_ =  zeta2_ * zeta2_
            fb_[i] = 1.0 - 0.5*zeta0_ + ZETA1_12 * zeta2_ - ZETA1_720 * zeta4_
def bernoulli_function_fast(zeta):
    fb_ = np.zeros_like(zeta)
    bernoulli_function_fill(zeta, fb_)
    return fb_

注意:如果您使用numba的新版本，您可以将这两者合并成相同的函数。

在我的机器上：

#create some test data
zeta = random.uniform(-1,1, size=2**24)
zeta[random.choice(len(zeta),size=2**23,replace=False )] = EPSILONZETA/2
>>> alltrue(bernoulli_function_fast(zeta)==bernoulli_function1(zeta))
True
>>> %timeit bernoulli_function1(zeta) # your function
1 loops, best of 3: 1.49 s per loop
>>> %timeit bernoulli_function_fast(zeta) #numba function
1 loops, best of 3: 347 ms per loop

这是~4倍的速度和更容易阅读。