Блог

Мне нужно выполнить вычисления вдоль оси массива NumPy, и я хотел бы знать, есть ли какой-либо способ обойти цикл for, поскольку это кажется очень медленным. Обычно я всегда пытаюсь векторизовать свой код, но с этим вычислением я просто не знаю как.

Итак, допустим, у меня есть 2d-массив X с формой (7000 x 10). Примерно половина записей — это значения NaN. Теперь я просматриваю каждую строку X (с именем x_i), которая представляет собой одномерный массив с формой (10). Основываясь на индексах значений, отличных от NaN в x_i, я разделяю 1d-массив M, форма (10), и 2d-массив C, форма (10 x 10). И затем я выполняю некоторые вычисления с ним и повторяю все вдоль первой оси X.

Мой код (довольно упрощенный) выглядит более или менее так:

 X = np.random.random_sample((7000, 10))  # Note that in the real case, X also has a lot of NaN values
M = np.random.random_sample(10)
C = np.random.random_sample((10, 10))

Res = np.empty(X.shape[0])   # Preallocation of result

for i in range(X.shape[0]):
    # Get row i of X
    x_i = X[i]

    # Get indices of the non-NaN values in x_i
    index_not_nan_i = np.where(~np.isnan(x_i))[0]
    
    # Partition M and C according to indices
    M_i = M[index_not_nan_i]
    C_i = C[index_not_nan_i[..., None], index_not_nan_i]


    Res[i] = M_i @ C_i @ M_i
  

Пример:

Допустим, первая строка X выглядит следующим образом x_i = [5, 3, 6, NaN, 7, NaN, NaN, 0, 5, NaN]; Таким образом, index_not_nan_i = [0,1,2,4,7,8]. Поскольку в x_i есть 6 значений, отличных от NaN, M_i будет иметь форму (6), а C_i будет иметь форму (6,6). Следующий x_i, конечно, будет иметь NaN-значения в других позициях, поэтому index_not_nan_i, возможно, будет [3,7,5], а M_i и C_i будут иметь форму (3) и (3 x 3).

Существуют ли какие-либо возможности векторизации? Я уже пробовал numpy.apply_along_axis(), но это, похоже, обычный цикл for с точки зрения производительности. Также, согласно документации, numpy.vectorize() также кажется, что это цикл for, хотя я еще не пробовал.