Search code examples
pythonnumpycomplex-numbers

How do I get the imaginary part (or real part) of an array in Python?


I have trouble with my N qubit wavefunction. I try to prepare a state psi = a|0> + b*exp(2i\pi\theta)|1>, and I wanted to check if the values for b*exp(2i\pi\theta) were well distributed.

Here is how I got my wavefunction :

N = 100
psi = np.full([100, 2], None)
for i in range(N) :
x = np.random.random()
y = y = np.exp(2j*np.pi*np.random.random())*np.sqrt(1-x**2)
psi[i] = [x,y]

I then used this line to get an array with only the y's and tried to plot them on the complex plane :

psi2 = psi[:,1]
plt.plot(psi2.real,psi2.imag)

I can't grasp why it doesn't plot the imaginary part and I just get :

Result of psi2 in the complex plane


Solution

  • You need to make the output array psi "complex aware". An easy way is to fill it with complex values instead of None objects:

    psi = np.full([100, 2], None)
    print(psi.dtype)
    # object  ... not good
    
    psi = np.full([100, 2], 0j)
    print(psi.dtype)
    # complex128   ... numpy inferred the complex data type for psi
    

    Now the .real and .imag attributes should work as expected.

    plt.plot(psi2[:,1].real,psi2[:,1].imag, '.')
    

    enter image description here