Pregunta

What's the easiest way to get the DFT matrix for 2-d DFT in python? I could not find such function in numpy.fft. Thanks!

¿Fue útil?

Solución 2

I don't think this is built in. However, direct calculation is straightforward:

import numpy as np
def DFT_matrix(N):
    i, j = np.meshgrid(np.arange(N), np.arange(N))
    omega = np.exp( - 2 * pi * 1J / N )
    W = np.power( omega, i * j ) / sqrt(N)
    return W

EDIT For a 2D FFT matrix, you can use the following:

x = np.zeros(N, N) # x is any input data with those dimensions
W = DFT_matrix(N)
dft_of_x = W.dot(x).dot(W)

Otros consejos

The easiest and most likely the fastest method would be using fft from SciPy.

import scipy as sp

def dftmtx(N):
    return sp.fft(sp.eye(N))

If you know even faster way (might be more complicated) I'd appreciate your input.

Just to make it more relevant to the main question - you can also do it with numpy:

import numpy as np

dftmtx = np.fft.fft(np.eye(N))

When I had benchmarked both of them I have an impression scipy one was marginally faster but I have not done it thoroughly and it was sometime ago so don't take my word for it.

Here's pretty good source on FFT implementations in python: http://nbviewer.ipython.org/url/jakevdp.github.io/downloads/notebooks/UnderstandingTheFFT.ipynb It's rather from speed perspective, but in this case we can actually see that sometimes it comes with simplicity too.

As of scipy 0.14 there is a built-in scipy.linalg.dft:

Example with 16 point DFT matrix:

>>> import scipy.linalg
>>> import numpy as np
>>> m = scipy.linalg.dft(16)

Validate unitary property, note matrix is unscaled thus 16*np.eye(16):

>>> np.allclose(np.abs(np.dot( m.conj().T, m )), 16*np.eye(16))
True

For 2D DFT matrix, it's just a issue of tensor product, or specially, Kronecker Product in this case, as we are dealing with matrix algebra.

>>> m2 = np.kron(m, m) # 256x256 matrix, flattened from (16,16,16,16) tensor

Now we can give it a tiled visualization, it's done by rearranging each row into a square block

>>> import matplotlib.pyplot as plt
>>> m2tiled = m2.reshape((16,)*4).transpose(0,2,1,3).reshape((256,256))
>>> plt.subplot(121)
>>> plt.imshow(np.real(m2tiled), cmap='gray', interpolation='nearest')
>>> plt.subplot(122)
>>> plt.imshow(np.imag(m2tiled), cmap='gray', interpolation='nearest')
>>> plt.show()

Result (real and imag part separately):

2D DFT basis

As you can see they are 2D DFT basis functions

Link to documentation

@Alex| is basically correct, I add here the version I used for 2-d DFT:

def DFT_matrix_2d(N):
    i, j = np.meshgrid(np.arange(N), np.arange(N))
    A=np.multiply.outer(i.flatten(), i.flatten())
    B=np.multiply.outer(j.flatten(), j.flatten())
    omega = np.exp(-2*np.pi*1J/N)
    W = np.power(omega, A+B)/N
    return W

Lambda functions work too:

dftmtx = lambda N: np.fft.fft(np.eye(N))

You can call it by using dftmtx(N). Example:

In [62]: dftmtx(2)
Out[62]: 
array([[ 1.+0.j,  1.+0.j],
       [ 1.+0.j, -1.+0.j]])

If you wish to compute the 2D DFT as a single matrix operation, it is necessary to unravel the matrix X on which you wish to compute the DFT into a vector, as each output of the DFT has a sum over every index in the input, and a single square matrix multiplication does not have this ability. Taking care to be sure we are handling the indices correctly, I find the following works:

M = 16
N = 16
X = np.random.random((M,N)) + 1j*np.random.random((M,N))
Y = np.fft.fft2(X)
W = np.zeros((M*N,M*N),dtype=np.complex)
​
hold = []
for m in range(M):
    for n in range(N):
        hold.append((m,n))
​
for j in range(M*N):
    for i in range(M*N):
        k,l = hold[j]
        m,n = hold[i]
        W[j,i] = np.exp(-2*np.pi*1j*(m*k/M + n*l/N))

np.allclose(np.dot(W,X.ravel()),Y.ravel())
True

If you wish to change the normalization to orthogonal, you can divide by 1/sqrt(MN) or if you wish to have the inverse transformation, just change the sign in the exponent.

This might be a little late, but there is a better alternative for creating the DFT matrix, that performs faster, using NumPy's vander

also, this implementation does not use loops (explicitly)

def dft_matrix(signal):

    N = signal.shape[0]  # num of samples
    w = np.exp((-2 * np.pi * 1j) / N)  # remove the '-' for inverse fourier
    r = np.arange(N)
    w_matrix = np.vander(w ** r, increasing=True)  # faster than meshgrid
    return w_matrix

if I'm not mistaken, the main improvement is that this method generates the elements of the power from the (already calculated) previous elements


you can read about vander in the documentation: numpy.vander

Licenciado bajo: CC-BY-SA con atribución
No afiliado a StackOverflow
scroll top