matplotlib 画3D的bug

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

x_= [0.67674598+0.j ,0.59789998+0.j, 0.61108673+0.j , 0.76330179+0.j,
 0.62970911+0.j]
x=np.array(x_)
print(type(x_[0]),x.dtype) #<class 'complex'> complex128

cmap = plt.cm.winter
cmp(x) #报错

错误信息

rgba = lut[xa]
IndexError: arrays used as indices must be of integer (or boolean) type

看了下.../python3.7/site-packages/matplotlib/colors.py下的__call__函数。
原因出在不能处理复数(complex)的映射。

逻辑是如果输入的x是浮点数，就 xa = x*self.N，取一个最近的整数值。
如果输入是整数，就跳过。
然后对xa的范围夹紧，使之不超过self.N的可表示范围。
最后返回 self._lut[xa]。
其中self._lut是一个形状为[self.N,4]的np.ndarray，存放的4维数值是(r,g,b,alpha)。

问题在于，plot 3D图形时，传入的X很可能是复数类型的。
Colormap类没有处理复数的逻辑。

虽然我把xa打印出来了，可以看到xa的虚部都是0。

xa= [0.67674598+0.j 0.59789998+0.j 0.61108673+0.j 0.76330179+0.j
0.62970911+0.j]

但是__call__里面只有一行

if xa.dtype.kind == “f”:

用来判断是不是f开头的浮点数。

解决bug的方法也简单，
在/site-packages/matplotlib/colors.py的__call__函数里，
加一段对"c"的判断，这样就能去掉虚部。剩下的继续按float流程走。

        if xa.dtype.kind == "c":
            xa = np.real(xa)
        if xa.dtype.kind == "f":
            with np.errstate(invalid="ignore"):
                xa *= self.N
                # Negative values are out of range, but astype(int) would
                # truncate them towards zero.
                xa[xa < 0] = -1
                # xa == 1 (== N after multiplication) is not out of range.
                xa[xa == self.N] = self.N - 1
                # Avoid converting large positive values to negative integers.
                np.clip(xa, -1, self.N, out=xa)
                xa = xa.astype(int)

想提交bugfix。