如何在numpy中向量化包含if语句的for循环？

Question

我对 numpy 还很陌生。我有以下代码，它给了我想要的结果：

import numpy as np

def get_result(S,K,delS):
    res=np.zeros(S.shape[0],dtype=np.float64)
    for j in range(res.shape[0]):
        if S[j]-K>delS:
            res[j]+=np.floor((S[j]-K)/delS)
            K+=np.floor((S[j]-K)/delS)*delS
        elif S[j]-K<-delS:
            res[j]+=np.ceil((S[j]-K)/delS)
            K+=np.ceil((S[j]-K)/delS)*delS
    return res

S=np.array([1.0,1.05,1.1,1.12,1.09,1.14,1.21,1.6,1.05,1.0,0.95,0.90,0.87,0.77,0.63,0.85,0.91,0.76],dtype=np.float64)

K=1.0
delS=0.1

l=get_result(S,K,delS)

for j in range(S.shape[0]):
    print("%d\t%.2f\t%.0f" % (j,S[j],l[j]))

然而，get_result 函数包含一个 for 循环，因此对于较大的输入向量 S 来说速度很慢。这样的函数可以用 numpy 语法向量化吗？任何帮助将不胜感激。

Answer 1

在处理数组或计算中的 2 个或多个条件时，一般模式是构造一个布尔掩码，并对掩码为 true 的元素执行一个操作，而在掩码为 false 的情况下执行不同的操作：

res=np.zeros(S.shape[0],dtype=np.float64)
mask - S-K>delS
res[mask] = ...S[mask]
res[~mask] = ...S[~mask]

一个变体是用np.where(mask)来识别索引。

但使您的计算复杂化的是测试不断变化。也就是说，j+1 的 K 源自 j 的计算。

for j in range(res.shape[0]):
    if S[j]-K>delS:
        res[j]+=np.floor((S[j]-K)/delS)
        K+=np.floor((S[j]-K)/delS)*delS
    elif S[j]-K<-delS:
        res[j]+=np.ceil((S[j]-K)/delS)
        K+=np.ceil((S[j]-K)/delS)*delS

为了处理这种迭代，我们通常尝试使用np.cumsum 或其他累积的ufunc 方法。

作为一般规则，numpy 代码在计算可以“并行”应用于数组（或数组集）的所有元素时是最快和最简单的 - 也就是说，以不依赖于迭代的顺序。然后我们可以将操作委托给快速编译的 numpy 函数，如数组加法和乘法。如果计算本质上是串行的（j 的值取决于 j-1），这就变得更棘手了。

所以，如果我的粗略阅读是正确的，那么困难的不是 if 语句，而是循环的串行性质。

===========================

玩弄，我能够删除if（实际上有3个子案例），但它仍然必须是迭代的：

def get_result(S,K,delS):
    S = S/delS
    res=np.zeros(S.shape[0],dtype=np.float64)
    for j in range(res.shape[0]):
        x = S[j] - K/delS
        xa = np.floor(np.abs(x)) * np.sign(x)
        res[j] += xa
        K += xa*delS         
    return res