我有这个简单的代码,试图获取两个复数E1
和E2
的实部的3D图,作为t
和g
的函数。
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import cmath
eps=0.5
def ReE1(t,g):
E1=eps+cmath.sqrt(t**2-g**2)
return E1.real
def ReE2(t,g):
E2=eps-cmath.sqrt(t**2-g**2)
return E2.real
fig = plt.figure()
ax = plt.axes(projection="3d")
t = np.linspace(0, 10, 50)
g = np.linspace(0, 10, 50)
X, Y = np.meshgrid(t, g)
Z = ReE1(X, Y)
ax.plot_surface(X, Y, Z, rstride=1, cstride=1,
cmap='winter', edgecolor='none')
Z = ReE2(X, Y)
ax.plot_surface(X, Y, Z, rstride=1, cstride=1,
cmap='summer', edgecolor='none')
plt.show()
在使用Python 3运行时出现以下错误。
Traceback (most recent call last):
File "t2.py", line 28, in <module>
Z = ReE1(X, Y)
File "t2.py", line 11, in ReE1
E1=eps+cmath.sqrt(t**2-g**2)
TypeError: only length-1 arrays can be converted to Python scalars
我们如何解决?另外,我们可以直接使用复杂函数E1
和E2
(而不是ReE1
和ReE2
)并在绘制时调用real
模块吗?
问题似乎是sqrt
中的cmath
仅接受标量,而您试图通过为其提供2维数组以向量化的方式使用它。一种解决方案是通过如下循环遍历cmath.sqrt
和t
的每个元素来应用g
:
def ReE1(t,g):
E1 = np.zeros(t.shape, dtype='complex')
for i in range(t.shape[0]):
for j in range(t.shape[1]):
E1[i][j]=eps+cmath.sqrt(t[i][j]**2-g[i][j]**2)
return E1.real
def ReE2(t,g):
E2 = np.zeros(t.shape, dtype='complex')
for i in range(t.shape[0]):
for j in range(t.shape[1]):
E2[i][j]=eps-cmath.sqrt(t[i][j]**2-g[i][j]**2)
return E2.real