Применить функцию к каждой матрице массива

Question

Я хотел бы применить функцию к каждой из матриц 3x3 в моем массиве (6890,6890,3,3). До сих пор я пытался использовать векторизацию на меньшем примере и с более простой функцией, которая не сработала.

def myfunc(x):
    return np.linalg.norm(x)

m = np.arange(45).reshape(5,3,3)
t = m.shape[0]
r = np.zeros((t, t))

q = m[:,None,...] @ m.swapaxes(1,2) # m[i] @ m[j].T
f = np.vectorize(q, otypes=[np.float])
res = myfunc(f)

Является ли векторизация правильным подходом для эффективного решения этой проблемы или мне следует попробовать что-то еще? Я также изучил numpy.apply_along_axis, но это относится только к 1D-подмассивам.

Answer 1

Вам нужно зациклить каждый элемент и применить функцию:

import numpy as np

# setup function
def myfunc(x):
    return np.linalg.norm(x*2)

# setup data array
data = np.arange(45).reshape(5, 3, 3)

# loop over elements and update
for item in np.nditer(data, op_flags = ['readwrite']):
    item[...] = myfunc(item)

Если вам нужно применить функцию для всего массива 3x3, используйте:

out_data = []
for item in data:
    out_data.append(myfunc(item))

Вывод:

[14.2828568570857, 39.761790704142086, 66.4529909033446, 93.32202312423365, 120.24974012445931]