Условная сумма по матрицам в python / numpy

Question

У меня есть два пустых массива X и W, каждый из которых имеет форму (N,N), полученную в конце расчета. Разделите диапазон X на равные интервалы [min(X), min(X)+delta, min(X)+2*delta,..., max(X)]. Я хотел бы знать, учитывая начальную точку интервала v, сумму соответствующих W значений:

idx = (X>=v) & (X<(v+delta))
W[idx].sum()

Мне нужна эта сумма для всех начальных интервалов (т. Е. Всего диапазона X), и мне нужно сделать это для множества различных матриц X и W. Профилирование определило, что это узкое место. То, что я делаю сейчас, составляет:

W_total = []
for v0, v1 in zip(X, X[1:]):
    idx = (X>=x0) & (X<x1)
    W_total.append( W[idx].sum() )

Как я могу ускорить это?

Answer 1

Вы пробовали numpy.histogram?

nbins = (X.max() - X.min()) / delta
W_total = np.histogram(X, weights=W, bins=nbins)

Answer 2

Вы можете использовать numpy.histogram() для вычисления всех этих сумм за одну операцию:

sums, bins = numpy.histogram(
    X, bins=numpy.arange(X.min(), X.max(), delta), weights=W)