Complex number troubles with numpy

Question

I'm attempting to translate some matlab code again and I've run into another pickle. The code itself is very simple, it's just a demonstration of a 4 node twiddle factor. Here is my attempt:

from numpy import *
from matplotlib import pyplot as plt 


x = zeros(4)
x[-1+1] = 0
x[0+1] = 1
x[1+1] = 1
x[2+1] = 0

z = 0 - 1j
W4 = exp(z*2*pi/4)
W0 = W4 ** 0
W1 = W4 ** 1
W2 = W4 ** 2
W3 = W4 ** 3


X = zeros(4)
X[-1+1] = (x[-1+1] + x[1+1]*W0) + W0*(x[0+1] + x[2+1]*W0)
X[0+1] = (x[-1+1] + x[1+1]*W2) + W1*(x[0+1] + x[2+1]*W2)
X[1+1] = (x[-1+1] + x[1+1]*W0) + W2*(x[0+1] + x[2+1]*W0)
X[2+1] = (x[-1+1] + x[1+1]*W2) + W3*(x[0+1] + x[2+1]*W2)


fx = fft.fft(x) 

plt.plot(X)
plt.plot(fx, 'ro')
plt.title("Results 4-point hand programmed FFT (blue) and the PYTHON routine (red o)")
plt.show()

Here are the output images. The first one is run with (almost) identical matlab code, the second one is the image from the python code above.

For lines 24 to 27 it gives me the error "ComplexWarning: Casting complex values to real discards the imaginary part". Now I'm not used to working with complex numbers in python. I tried adding a complex component to all the variables, but it gave me a graph that's way off from the matlab one. Thoughts? If you would like me to post the matlab code as well, let me know.

Answer 1

When you specify the array x and X you need to make sure it is of complex data type, i.e:

x = zeros((4),dtype=complex)

EDIT:

To fix the plot you need to plot both real and imaginary parts:

plt.plot(X.real,X.imag)
plt.plot(fx.real,fx.imag, 'ro')

This gives me:

....which looks like your Matlab graph.