Многочисленные арифметические операции эффективно на 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() предполагается запускать много раз в моем проекте, поэтому я хотел бы переписать его более питоническим образом, чтобы сделать его максимально эффективным.Я попытался с помощью следующего метода:

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 для вычисления времени.)