Python 线框图

该程序根据50层（d1、d2）的厚度计算多层系统的反射系数。如果我取任意两个数字（d1，d2），它就正确了。但我需要得到线框图，其中d1，d2在某个范围内有意义，我在第13行得到一个错误：“ValueError:input必须是一个正方形数组”。我怎样才能修好它

from math import pi
import numpy as np
import matplotlib.pyplot as plt

def R(n1, n2, d1, d2, lamda):
    phy1 = (-2*pi*n1*d1/lamda)
    phy2 = (-2*pi*n2*d2/lamda)
    DPD1 = 0.5*np.array([[2*np.cos(phy1),   2j*np.sin(phy1)/n1],    [n1*2j*np.sin(phy1),  2*np.cos(phy1) ]])
    DPD2 = 0.5*np.array([[2*np.cos(phy2),   2j*np.sin(phy2)/n2],   [n2*2j*np.sin(phy2),  2*np. cos(phy2) ]])
    D0 = 0.5 * np.array([[1, 1], [1, -1]])
    DS = np.array([[1, 1], [n1, -n1]])
    DPD = np.dot(DPD1, DPD2)
    DPD = np.linalg.matrix_power(DPD, 50)
    M = np.dot(D0, DPD)
    M = np.dot(M, DS)
    return(abs(M[1,0]/M[0,0])**2)

x = np.arange(0, 10, 1)
y = np.arange(0, 10, 1)
X, Y = np.meshgrid(x, y)
Z = R(0.99910053+0.00183184j,  0.92373900+0.00644652j, X, Y, 13.5)
fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_wireframe(X, Y, Z, antialiased=True)

---

我不知道没有for循环怎么做。以下是我的解决方案：

from math import pi
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def R(n1, n2, d1, d2, lamda):
    phy1 = (-2*pi*n1*d1/lamda)
    phy2 = (-2*pi*n2*d2/lamda)
    DPD1 = 0.5*np.array([[2*np.cos(phy1),   2j*np.sin(phy1)/n1],    [n1*2j*np.sin(phy1),  2*np.cos(phy1) ]])
    DPD2 = 0.5*np.array([[2*np.cos(phy2),   2j*np.sin(phy2)/n2],   [n2*2j*np.sin(phy2),  2*np. cos(phy2) ]])
    DPD = np.dot(DPD1, DPD2)
    DPD = np.linalg.matrix_power(DPD, 50)
    D0 = 0.5 * np.array([[1, 1], [1, -1]])
    DS = np.array([[1, 1], [n1, -n1]])
    M = np.dot(D0, DPD)
    M = np.dot(M, DS)
    return(abs(M[1,0]/M[0,0])**2)

x = np.arange(0, 10, 1)
y = np.arange(0, 10, 1)
X, Y = np.meshgrid(x, y)
Z = np.zeros(X.shape).ravel()
for i, (x, y) in enumerate(zip(X.ravel(), Y.ravel())):
    Z[i] = R(0.99910053+0.00183184j,  0.92373900+0.00644652j, x, y, 13.5)
fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_wireframe(X, Y, Z.reshape(X.shape), antialiased=True)
plt.show()

---

我不知道没有for循环怎么做。以下是我的解决方案：

from math import pi
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def R(n1, n2, d1, d2, lamda):
    phy1 = (-2*pi*n1*d1/lamda)
    phy2 = (-2*pi*n2*d2/lamda)
    DPD1 = 0.5*np.array([[2*np.cos(phy1),   2j*np.sin(phy1)/n1],    [n1*2j*np.sin(phy1),  2*np.cos(phy1) ]])
    DPD2 = 0.5*np.array([[2*np.cos(phy2),   2j*np.sin(phy2)/n2],   [n2*2j*np.sin(phy2),  2*np. cos(phy2) ]])
    DPD = np.dot(DPD1, DPD2)
    DPD = np.linalg.matrix_power(DPD, 50)
    D0 = 0.5 * np.array([[1, 1], [1, -1]])
    DS = np.array([[1, 1], [n1, -n1]])
    M = np.dot(D0, DPD)
    M = np.dot(M, DS)
    return(abs(M[1,0]/M[0,0])**2)

x = np.arange(0, 10, 1)
y = np.arange(0, 10, 1)
X, Y = np.meshgrid(x, y)
Z = np.zeros(X.shape).ravel()
for i, (x, y) in enumerate(zip(X.ravel(), Y.ravel())):
    Z[i] = R(0.99910053+0.00183184j,  0.92373900+0.00644652j, x, y, 13.5)
fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_wireframe(X, Y, Z.reshape(X.shape), antialiased=True)
plt.show()