Using functions to create Electric Field array for 2d Density Plot and 3d Surface Plot

Question

Below is my code, I'm supposed to use the electric field equation and the given variables to create a density plot and surface plot of the equation. I'm getting "invalid dimensions for image data" probably because the function E takes multiple variables and is trying to display them all as multiple dimensions. I know the issue is that I have to turn E into an array so that the density plot can be displayed, but I cannot figure out how to do so. Please help.

import numpy as np
from numpy import array,empty,linspace,exp,cos,sqrt,pi
import matplotlib.pyplot as plt

lam = 500 #Nanometers
x = linspace(-10*lam,10*lam,10)
z = linspace(-20*lam,20*lam,10)
w0 = lam
E0 = 5

def E(E0,w0,x,z,lam):
    E = np.zeros((len(x),len(z)))
    for i in z:
        for j in x:
            E = ((E0 * w0) / w(z,w0,zR(w0,lam)))
            E = E * exp((-r(x)**2) / (w(z,w0,zR(w0,lam)))**2)
            E = E * cos((2 * pi / lam) * (z + (r(x)**2 / (2 * Rz(z,zR,lam)))))
    return E

def r(x):
    r = sqrt(x**2)
    return r

def w(z,w0,lam):
    w = w0 * sqrt(1 + (z / zR(w0,lam))**2)
    return w

def Rz(z,w0,lam):
    Rz = z * (1 + (zR(w0,lam) / z)**2)
    return Rz

def zR(w0,lam):
    zR = pi * lam
    return zR

p = E(E0,w0,x,z,lam)

plt.imshow(p)

Answer 1

It took me way too much time and thinking but I finally figured it out after searching for similar examples of codes for similar problems. The correct code looks like:

import numpy as np
from numpy import array,empty,linspace,exp,cos,sqrt,pi
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

lam = 500*10**-9 #Nanometers
x1 = linspace(-10*lam,10*lam,100)
z1 = linspace(-20*lam,20*lam,100)

[x,y] = np.meshgrid(x1,z1)
w0 = lam
E0 = 5
r = sqrt(x**2)
zR = pi * lam
w = w0 * sqrt(1 + (y / zR)**2)
Rz = y * (1 + (zR / y)**2)

E = (E0 * w0) / w
E = E * exp((-r**2 / w**2))
E = E * cos((2 * pi / lam) * (y + (r**2 / (2 * Rz))))

def field(x,y):
    lam = 500*10**-9
    k = (5 * lam) / lam * sqrt(1 + (y / (pi*lam))**2)
    k *= exp(((-sqrt(x**2)**2 / (lam * sqrt(1 + (y / pi * lam)**2))**2)))
    k *= cos((2 / lam) * (y + ((sqrt(x**2)**2 / (2 * y * (1 + (pi * lam / y)**2))))))
    return k

#Density Plot
f = field(x,y)
plt.imshow(f)
plt.show()

#Surface Plot
fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_surface(x,y,E,rstride=1,cstride=1)
plt.show