Ohm
Ohm

Reputation: 2442

Python: numpy fftn over a list of numpy array

I am trying to efficiently np.fft.fftn and array of 2D numpy arrays. V0 is an array of shape (nvar,nx,ny), and I would like to perform FFT over each 2D array from the first dimension of V0. What I tried to do here is to compare between iteration over the first axis of V0, and compare it with a naive attempt to do a FFT over the whole multidimensional array:

In [45]: import numpy as np

In [46]: V0 = np.random.random((3,128,128))

In [47]: V0fft = []

In [48]: for i in xrange(V0.shape[0]):
    ...:     V0fft.append(np.fft.fftn(V0[i]))
    ...:     

In [49]: V0fftdirect = np.fft.fftn(V0)

In [50]: np.amax(np.abs(V0fft - V0fftdirect))
Out[50]: 16366.207818488827

So how can I do it correctly?

Upvotes: 1

Views: 778

Answers (2)

Chr
Chr

Reputation: 965

  • Replace np.fft.fftn by np.fft.fft2 in the for loop.
  • Specify the fftn axis : numpy.fft.fftn(a, s=None, axes=(-2, -1), norm=None)

Upvotes: 1

Ohm
Ohm

Reputation: 2442

V0fftdirect = np.fft.fft2(V0) did the job :

In [93]: import numpy as np

In [94]: V0 = np.random.random((3,128,128))

In [95]: V0fftfor = []

In [96]: for i in xrange(V0.shape[0]):
    ...:     V0fftfor.append(np.fft.fftn(V0[i]))
    ...:     

In [97]: V0fftdirect = np.fft.fft2(V0)

In [98]: np.amax(np.abs(V0fftfor - V0fftdirect))
Out[98]: 0.0

Upvotes: 1

Related Questions