在NumPy上高效地执行多种算术运算

Question

对于Python程序，我具有以下数据。

import numpy as np

np.random.seed(28)

n = 100000
d = 60

S = np.random.rand(n)
O = np.random.rand(n, d, d)
p = np.random.rand()
mask = np.where(S < 0.5)

我想运行以下算法：

def method1():
    sum_p = np.zeros([d, d])
    sum_m = np.zeros([d, d])

    for k in range(n):
        s = S[k] * O[k]
        sum_p += s
        if(S[k] < 0.5):
            sum_m -= s

    return p * sum_p + sum_m

这是一个最小的示例，但是method1()中的代码应该在我的项目中运行多次，因此我想以更加Python化的方式重写它，以使其尽可能高效。我尝试了以下方法：

def method2():
    sall = S[:, None, None] * O
    return p * sall.sum(axis=0) - sall[mask].sum(axis=0)

但是，尽管此方法在d较低的情况下效果更好，但在d=60时却无法提供良好的时机：

# To check that both methods provide the same result.
In [1]: np.sum(method1() == method2()) == d*d 
Out[1]: True

In [2]: %timeit method1()
Out[2]: 801 ms ± 2.98 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [3]: %timeit method2()
Out[3]: 1.91 s ± 6.17 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

您还有其他想法可以优化此方法吗？

（作为附加信息，应该在最终代码的其他部分中使用变量mask，因此我无需在method2的代码中考虑它来进行时间计算）