使用Matplotlib在3D绘图中使用cmath的问题

Issues with cmath in 3D plotting using Matplotlib

来源: link 标签:pythonpython-3.xmatplotlib3dcmath 更新时间:

  问题

我有这个简单的代码,它试图获得两个复杂数字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只接受scalars,而您正在尝试通过提供2-d数组以矢量化方式使用它.一个解决方案是通过循环遍历t和g的每个元素上应用cmath.sqrt,如下所示: 

 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

enter image description here

  相同标签的其他问题

pythonpython-3.xmatplotlib3dcmath
上一个问题:
下一个问题: