Use numpy for FFT convolution

Question

I have two arrays, a and b both of length N. I want to do, for example j = 4

import numpy as np
N = 1000
c = np.zeros(N)
j = 40
c[j] = np.sum([a[i] * b[(j-i)%(N-1)] for i in range(N)])

But if directly using convolution theorem, I would get

c[j] = np.sum([a[i] * b[j-i] for i in range(N)])

which is the same as

c = np.fft.ifft(np.fft(a) * np.fft(b))

Is there a clever way to pad zeros such that one index misalignment can be treated using np.fft, which is very fast.

Answer 1

The solution is to first truncate the arrays so that the rolling back in numpy.arrays, e.g., index -1 in the truncated arrays is the second last in the original one. For example,

import numpy as np
N = 100
a_arr = np.random.rand(N)
b_arr = np.random.rand(N)
c_arr = np.zeros(N)

for j in range(N):
    c_arr[j] = np.sum([a_arr[i] * b_arr[(j-i) % (N-1)] for i in range(N)])

# Truncated to the first N-1 elements.
a_arr_t = a_arr[:N-1]
b_arr_t = b_arr[:N-1]

# Use FFT for convolution for the truncated arrays. 
# So that rolling back -1 means the second last not the last.
c_arr_t = np.fft.ifft(np.fft.fft(a_arr_t) * np.fft.fft(b_arr_t))
for j in range(N-1):
    c_arr_t[j] += a_arr[N-1] * b_arr[(j-(N-1)) % (N-1)]

# The last element is treated separately.
c_arr_check = np.append(c_arr_t, np.sum([a_arr[i] * b_arr[(N-1-i) % (N-1)] for i in range(N)]))
print(np.max(c_arr_check - c_arr))

And finally add the tail element manually.

The total complexity for this is still O(Nlog(N)) plus some O(N) operations to complement tails.