使用 numpy 对带有条件代码的循环进行向量化

Question

我正在尝试将简单代码转换为矢量化器。我是这个概念的新手。

我尝试过使用np.vectorize 选项，但它会引发错误：

你能帮我解决这个问题吗？有没有其他更好的方法来矢量化它？

---

代码：

   def test(s,m):
      mean_matrix=s
      m2=s/m // forming an array
      result=[]
      for x in m2:
         return x/5



  s=np.random.random(50)
  m=np.random.random(100)
  temp=np.vectorize(test,otypes=[np.float],cache=False)
  out = temp(s,m)

Answer 1

首先将您的代码分成循环中使用的部分（可以矢量化）和其余部分（已经矢量化）：

def kernel(s, x):
    i = 0
    initial_value = s[0].copy()
    while initial_value < x:
        i = i+1
        initial_value += s[i]
    return i

def test_split(s, m):
    mean_matrix = s.mean() # 5 usec
    product_m2 = mean_matrix * m # forming an array, 2 usec
    return [kernel(s, x) for x in product_m2] # 167 usec

test_split 和test 做同样的事情，只是被重构了。我已经对每一行进行了注释，说明在那里花费了多少时间，以表明加快循环速度会有所作为。

如果你使用np.vectorize()，它看起来像这样：

def test_vectorize(s, m):
    mean_matrix = s.mean() # 5 usec
    product_m2 = mean_matrix * m # forming an array, 2 usec
    kernel_vec = np.vectorize(lambda x: kernel(s, x)) # 2 usec
    return kernel_vec(product_m2) # 194 usec

这更慢！那是因为np.vectorize() 不会让事情变得更快，它只会改变美学。它仍然在内部使用 Python 循环。最好忘记它的存在。

所以我们需要一种不同的方法，为此我们需要考虑kernel 做了什么。在我看来，它所做的是找到s 的最长初始子序列，其总和小于x。 s.cumsum() 将为我们提供我们需要的运行总数，但是我们如何在product_m2 中搜索小于每个x 的最大长度？我们不能使用循环，否则它会再次变慢。但是我们可以使用np.searchsorted()。这里是：

def test_fast(s, m):
     mean_matrix = s.mean() # 5 usec
     product_m2 = mean_matrix * m # forming an array, 2 usec
     s_sum = s.cumsum() # 1 usec
     return s_sum.searchsorted(product_m2) # 2 usec

这样我们就有了 14 倍的加速。