使用函数为2d密度图和3d表面图创建电场阵列

时间:2018-05-03 18:42:24

标签: python arrays function numpy matplotlib

下面是我的代码,我应该使用电场方程和给定的变量来创建方程的密度图和表面图。我得到"图像数据的尺寸无效"可能是因为函数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)

1 个答案:

答案 0 :(得分: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