numpy векторизация вместо циклов

Question

Я написал код в Python, который отлично работает, но очень медленный; Я думаю, из-за циклов. Надеюсь, вы можете ускорить выполнение следующих операций с помощью команд numpy. Позвольте мне определить цель.

Предположим, у меня есть двумерная матрицаразмеров row x col. Например, рассмотрите массив 6 x 11 (см. Рисунок ниже).

1. Я хочу, чтобы вычислить среднее значение для всех строк, то есть сумма ⱼ aᵢⱼ в результате чего в массиве. Это, конечно, можно легко сделать. (Я называю это значение CM_tilde)

2. Теперь, для каждая строка Я хочу, чтобы вычислить среднее значение некоторых выбранных значений, а именно все значения ниже определенного порога, вычисляя их сумму и деления на количество всех столбцов (N). Если значение превышает указанный порог, добавляется значение CM_tilde (среднее значение для всей строки). Это значение называется CM

3. Затем величина CM вычитается из каждого элемента в строке

В дополнение к этому я хочу иметь Numpy массив или список, в котором все эти CM значения перечислены ,

Фигура:

Следующий код работает, но очень медленно (особенно если массивы получать большие)

CM_tilde = np.mean(data, axis=1) 
N = data.shape[1] 
data_cm = np.zeros((data.shape[0], data.shape[1], data.shape[2])) 
all_CMs = np.zeros((data.shape[0], data.shape[2])) 
for frame in range(data.shape[2]): 
    for row in range(data.shape[0]): 
     CM=0 
     for col in range(data.shape[1]): 
      if data[row, col, frame] < (CM_tilde[row, frame]+threshold): 
       CM += data[row, col, frame] 
      else: 
       CM += CM_tilde[row, frame] 
     CM = CM/N 
     all_CMs[row, frame] = CM 
     # calculate CM corrected value 
     for col in range(data.shape[1]): 
      data_cm[row, col, frame] = data[row, col, frame] - CM 
    print "frame: ", frame 
return data_cm, all_CMs 

Есть идеи?

Answer 1

Это довольно легко векторизации, что вы делаете:

import numpy as np 

#generate dummy data 
nrows=6 
ncols=11 
nframes=3 
threshold=0.3 
data=np.random.rand(nrows,ncols,nframes) 

CM_tilde = np.mean(data, axis=1) 
N = data.shape[1] 

all_CMs2 = np.mean(np.where(data < (CM_tilde[:,None,:]+threshold),data,CM_tilde[:,None,:]),axis=1) 
data_cm2 = data - all_CMs2[:,None,:] 

Сравнивая это с оригиналов:

In [684]: (data_cm==data_cm2).all() 
Out[684]: True 

In [685]: (all_CMs==all_CMs2).all() 
Out[685]: True 

Логика в том, что мы работаем с массивами размером [nrows,ncols,nframes] одновременно. Основной трюк заключается в использовании радиовещания python, путем поворота CM_tilde размера [nrows,nframes] в CM_tilde[:,None,:] размера [nrows,1,nframes]. Затем Python будет использовать одни и те же значения для каждого столбца, поскольку это одноэлементный размер этого модифицированного CM_tilde.

С помощью np.where мы выбираем (на основе threshold), хотим ли мы, чтобы получить соответствующее значение data, или, опять же, стоимость вещания CM_tilde. Новое использование np.mean позволяет нам вычислить all_CMs2.

На последнем этапе мы использовали широковещательную передачу путем прямого вычитания этого нового all_CMs2 из соответствующих элементов data.

Это может помочь в векторизации кода таким образом, если посмотреть на неявные индексы ваших временных переменных. Я имею в виду, что ваша временная переменная CM живет внутри цикла над [nrows,nframes], и его значение сбрасывается с каждой итерацией. Это означает, что CM фактически является величиной CM[row,frame] (позже явно присвоенной 2-мерному массиву all_CMs), и отсюда легко увидеть, что вы можете построить его, суммируя соответствующее количество CMtmp[row,col,frames] по его размеру столбца. Если это поможет, вы можете назвать np.where(...) часть CMtmp для этой цели, а затем вычислить np.mean(CMtmp,axis=1). Тот же результат, очевидно, но, вероятно, более прозрачный.

Answer 2

Вот моя векторизация вашей функции. Я работал изнутри и прокомментировал более ранние версии, когда я пошел. Таким образом, первая петля, которую я векторизовал, имеет ### метки комментариев.

Это не так чисто и хорошо аргументировано, как @Andras's Ответ, но, надеюсь, это учебный материал, дающий представление о том, как вы можете обращаться к этой проблеме постепенно.

def foo2(data, threshold): 
    CM_tilde = np.mean(data, axis=1) 
    N = data.shape[1] 
    #data_cm = np.zeros((data.shape[0], data.shape[1], data.shape[2])) 
    ##all_CMs = np.zeros((data.shape[0], data.shape[2])) 
    bmask = data < (CM_tilde[:,None,:] + threshold) 
    CM = np.zeros_like(data) 
    CM[:] = CM_tilde[:,None,:] 
    CM[bmask] = data[bmask] 
    CM = CM.sum(axis=1) 
    CM = CM/N 
    all_CMs = CM.copy() 
    """ 
    for frame in range(data.shape[2]): 
     for row in range(data.shape[0]): 
      ###print(frame, row) 
      ###mask = data[row, :, frame] < (CM_tilde[row, frame]+threshold) 
      ###print(mask) 
      ##mask = bmask[row,:,frame] 
      ##CM = data[row, mask, frame].sum() 
      ##CM += (CM_tilde[row, frame]*(~mask)).sum() 

      ##CM = CM/N 
      ##all_CMs[row, frame] = CM 
      ## calculate CM corrected value 
      #for col in range(data.shape[1]): 
      # data_cm[row, col, frame] = data[row, col, frame] - CM[row,frame] 
     print "frame: ", frame 
    """ 
    data_cm = data - CM[:,None,:] 
    return data_cm, all_CMs 

Выходные совпадения для этого небольшого тестового корпуса, что более чем помогло мне получить размеры вправо.

threshold = .1 
data = np.arange(4*3*2,dtype=float).reshape(4,3,2)