python - 带比较运算符的 Numpy 广播；循环迭代

Question

我通过两种方式实现了循环迭代函数:

def Spin1(n, N) :           # n - current state, N - highest state
    value = n + 1   
    case1 = (value > N) 
    case2 = (value <= N)
    return case1 * 0 + case2 * value

def Spin2(n, N) :
    value = n + 1   
    if value > N :
        return 0
    else : return value

这些函数在返回结果方面是相同的。但是，第二个函数对于 numpy 数组不具有广播功能。因此，为了测试第一个函数，我运行以下命令:

import numpy

AR1 = numpy.zeros((3, 4), dtype = numpy.uint32)
AR1[1,2] = 5
print AR1
print Spin1(AR1,5)

它神奇地起作用了，这太甜蜜了。所以我清楚地看到了我想要的:

[[0 0 0 0]
 [0 0 5 0]
 [0 0 0 0]]
[[1 1 1 1]
 [1 1 0 1]
 [1 1 1 1]]

现在使用第二个函数print Spin2(AR1,5)，它失败并出现以下错误:

if value > N
ValueError: The truth value of an array with more than one element is ambiguous. 
Use a.any() or a.all()

原因很清楚，因为 if Array 语句是无稽之谈。所以现在我只使用第一个变体。但是，当我查看这些函数时，我有一种强烈的感觉，在第一个函数中存在更多的数学运算，因此我不会失去对它进行优化的希望。

问题:
1. 是否可以优化函数 Spin1 以减少操作或如何在广播模式下使用函数 Spin2(可能无需进行我的代码太丑了)？额外问题:对数组进行操作的最快方法是什么？
2.是否有一些标准的 Python 函数可以执行相同的计算(不隐式具有广播功能)以及如何正确调用它 - 可能是“循环增量”？

Answer 1

有一个 numpy 函数可以实现此目的:np.where :

In [590]: AR1
Out[590]: 
array([[0, 0, 0, 0],
       [0, 0, 5, 0],
       [0, 0, 0, 0]], dtype=uint32)

In [591]: np.where(AR1 >= 5, 0, 1)
Out[591]: 
array([[1, 1, 1, 1],
       [1, 1, 0, 1],
       [1, 1, 1, 1]])

<小时/>

所以，你可以定义:

def Spin1(n, N) :        
    value = n + 1   
    return np.where(value > N, 0, value)

<小时/>

NumPy 还提供了一种将普通 Python 函数转换为 ufuncs 的方法:

def Spin2(n, N) :
    value = n + 1   
    if value > N :
        return 0
    else : return value

Spin2 = np.vectorize(Spin2)

这样您现在就可以在数组上调用 Spin2:

In [595]: Spin2(AR1, 5)
Out[595]: 
array([[1, 1, 1, 1],
       [1, 1, 0, 1],
       [1, 1, 1, 1]])

然而，np.vectorize主要提供语法糖。仍然对每个数组元素进行 Python 函数调用，这使得 np.vectorized ufuncs no faster than equivalent code using Python for-loops .