Только свертка вдоль одной оси

У меня два двухмерных массива с одинаковыми размерами первой оси. В python я хотел бы сверлить только две матрицы по второй оси. Я хотел бы получить C ниже, не вычисляя свертку вдоль первой оси.

import numpy as np
import scipy.signal as sg

M, N, P = 4, 10, 20
A = np.random.randn(M, N)
B = np.random.randn(M, P)

C = sg.convolve(A, B, 'full')[(2*M-1)/2]

Есть ли быстрый способ?

Ответ 1

Вы можете использовать np.apply_along_axis для применения np.convolve вдоль нужной оси. Ниже приведен пример применения фильтра boxcar к 2d-массиву:

import numpy as np

a = np.arange(10)
a = np.vstack((a,a)).T

filt = np.ones(3)

np.apply_along_axis(lambda m: np.convolve(m, filt, mode='full'), axis=0, arr=a)

Это простой способ обобщить многие функции, не имеющие аргумента axis.

Ответ 2

С ndimage.convolve1d вы можете указать ось...

Ответ 3

np.apply_along_axis вам действительно не поможет, потому что вы пытаетесь выполнить итерацию по двум массивам. Фактически, вам нужно будет использовать цикл, как описано здесь.

Теперь петли прекрасны, если ваши массивы малы, но если N и P велики, то вы, вероятно, захотите использовать FFT для свертки.

Тем не менее, вам нужно соответствующим образом заполнить ваши массивы нулями, чтобы ваша "полная" свертка имела ожидаемую форму:

M, N, P = 4, 10, 20
A = np.random.randn(M, N)
B = np.random.randn(M, P)

A_ = np.zeros((M, N+P-1), dtype=A.dtype)
A_[:, :N] = A
B_ = np.zeros((M, N+P-1), dtype=B.dtype)
B_[:, :P] = B

A_fft = np.fft.fft(A_, axis=1)
B_fft = np.fft.fft(B_, axis=1)
C_fft = A_fft * B_fft

C = np.real(np.fft.ifft(C_fft))

# Test
C_test = np.zeros((M, N+P-1))
for i in range(M):
    C_test[i, :] = np.convolve(A[i, :], B[i, :], 'full')

assert np.allclose(C, C_test)

Ответ 4

для 2D-массивов функция scipy.signal.convolve2d работает быстрее, а scipy.signal.fftconvolve может быть еще быстрее (в зависимости от размеров массивов):

Здесь тот же код с N = 100000

import time    
import numpy as np
import scipy.signal as sg

M, N, P = 10, 100000, 20
A = np.random.randn(M, N)
B = np.random.randn(M, P)

T1 = time.time()
C = sg.convolve(A, B, 'full')
print(time.time()-T1)

T1 = time.time()
C_2d = sg.convolve2d(A, B, 'full')
print(time.time()-T1)

T1 = time.time()
C_fft = sg.fftconvolve(A, B, 'full')
print(time.time()-T1)

>>> 12.3
>>> 2.1
>>> 0.6

Ответы одинаковы с небольшими различиями из-за использования различных методов вычислений (например, fft против прямого умножения, но я не знаю, что использует exaclty convolve2d):

print(np.max(np.abs(C - C_2d)))
>>>7.81597009336e-14

print(np.max(np.abs(C - C_fft)))
>>>1.84741111298e-13