Ускоренная работа с 3D-матрицей - Python

Question

Я работаю с 3D-матрицей в Python, например, с такой матрицей размером 2x3x4:

[[[1 2 1 4]
  [3 2 1 1]
  [4 3 1 4]]

 [[2 1 3 3]
  [1 4 2 1]
  [3 2 3 3]]]

У меня есть задача найти значение энтропии в каждой строке в каждой матрице измерений.Например, в строке 1 размерности 1 матрицы выше [1,2,1,4] нормализованное значение (таким образом, общая сумма равна 1) равно [0.125, 0.25, 0.125, 0.5], а значение энтропии вычисляется по формуле -sum(i*log(i)), где i - этонормализованное значение.Результирующая матрица представляет собой матрицу 2x3, где в каждом измерении есть 3 значения энтропии (потому что есть 3 строки).

Вот рабочий пример моего кода, использующего случайную матрицу каждый раз:

from scipy.stats import entropy
import numpy as np

matrix = np.random.randint(low=1,high=5,size=(2,3,4)) #how if size is (200,50,1000)
entropy_matrix=np.zeros((matrix.shape[0],matrix.shape[1]))
for i in range(matrix.shape[0]):
    normalized = np.array([float(k)/np.sum(j) for j in matrix[i] for k in j]).reshape(matrix.shape[1],matrix.shape[2])
    entropy_matrix[i] = np.array([entropy(m) for m in normalized])

У меня вопрос, как мне масштабировать эту программу для работы с очень большой трехмерной матрицей (например, с размером 200x50x1000)?

Я использую Python в Windows 10 (с дистрибутивом Anaconda).Используя размер 3D-матрицы 200x50x1000, время работы моего компьютера составило 290 с.

Answer 1

Используя определение entropy для второй части и транслируемую операцию в первой части, одно векторизованное решение будет иметь вид -

p1 = matrix/matrix.sum(-1,keepdims=True).astype(float)
entropy_matrix_out = -np.sum(p1 * np.log(p1), axis=-1)

В качестве альтернативы, мы можем использовать einsum для второй части длядалее перф.буст -

entropy_matrix_out = -np.einsum('ijk,ijk->ij',p1,np.log(p1),optimize=True)