Использование функций для создания массива электрического поля для двухмерного графика плотности и трехмерного графика поверхности

Question

Ниже приведен мой код, я должен использовать уравнение электрического поля и заданные переменные для создания графика плотности и графика поверхности. Я получаю «недопустимые измерения для данных изображения», вероятно, потому что функция E принимает несколько переменных и пытается отобразить их все как несколько измерений. Я знаю, что проблема в том, что мне нужно превратить E в массив, чтобы можно было отобразить график плотности, но я не могу понять, как это сделать. Пожалуйста, помогите.

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

Это заняло у меня слишком много времени и размышлений, но я наконец понял это после поиска похожих примеров кодов для похожих проблем. Правильный код выглядит так:

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