ValueError: x and y must have same first dimension, but have shapes (10000,) and (1, 10000)-CodePudding

I am trying to write a code where I do a summation over a function with two variables, E and n, where n is a range of 0-5, and plot the result which should be a multiple step function.

import matplotlib.pyplot as plt
import numpy as np
 
E = np.linspace(0, 10, 10000)

for n in range(0, 6):
    h= []
    h.append(1/(1   np.exp(-2*np.pi * (E-(n-0.5)*3))))    
    print(h)
    

plt.plot(E,h)

This is the code I have been using. It produces multiple arrays for h:

[array([0.99991931, 0.99991981, 0.99992031, ..., 1.        , 1.        ,
       1.        ])]
[array([8.06930057e-05, 8.12016207e-05, 8.17134412e-05, ...,
       1.00000000e 00, 1.00000000e 00, 1.00000000e 00])]
[array([5.25548518e-13, 5.28861364e-13, 5.32195094e-13, ...,
       1.00000000e 00, 1.00000000e 00, 1.00000000e 00])]
[array([3.42258854e-21, 3.44416317e-21, 3.46587379e-21, ...,
       9.99999847e-01, 9.99999848e-01, 9.99999849e-01])]
[array([2.22893072e-29, 2.24298100e-29, 2.25711985e-29, ...,
       4.09276641e-02, 4.11750329e-02, 4.14238322e-02])]
[array([1.45157155e-37, 1.46072167e-37, 1.46992947e-37, ...,
       2.77912110e-10, 2.79663956e-10, 2.81426846e-10])]

but when I try to plot I get the following error:

ValueError: x and y must have same first dimension, but have shapes (10000,) and (1, 10000)

I don't undestand what is causing this and any help would be appreciated.

CodePudding user response：

I assume what you want is the following:

E = np.expand_dims(np.linspace(0, 10, 10000), 1)
n = np.arange(0, 6)

h = 1/(1   np.exp(-2*np.pi * (E-(n-0.5)*3)))

plt.plot(E, h)

Note that this vectorized calculation does not require a loop, as Numpy broadcasting figures out automatically how to combine the E and n vectors to a 2D h array, if the input dimension are consistent (note the np.expand_dims). In general, whenever you see the need to use a loop in Numpy, it's good advice to take a step back and think about vectorization.

Fixed version of your original code:

import matplotlib.pyplot as plt
import numpy as np
 
E = np.linspace(0, 10, 10000)

h= []
for n in range(0, 6):
    h.append(1/(1   np.exp(-2*np.pi * (E-(n-0.5)*3))).T)

    
plt.plot(E, np.stack(h).T)

Define h before the loop, such that it is not defined as the empty list on each iteration.
Combine the list of 1D arrays into a 6x10_000 2D array using np.stack and transpose (.T) in order to plot all curves.