3D plotでグラフの色を途中で変える

やりたいこと

こういうやつ

LineCollectionの3D版であるLine3DCollectionを使います.

連続変化版

Line3DCollectionに渡すcmapを変えて使用する配色セットを変えます.

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

theta  = np.linspace(0, 20 * np.pi, 500)
x = np.cos(theta)
y = np.sin(theta)
z = 0.1 * theta / np.pi
points = np.array([x, y, z]).T.reshape(-1, 1, 3)
segments = np.concatenate([points[:-1], points[1:]], axis=1)

fig = plt.figure()
ax = fig.gca(projection = '3d')
norm = plt.Normalize(z.min(),  z.max())
lc = Line3DCollection(segments, cmap='viridis', norm=norm)
lc.set_array(z)
lc.set_linewidth(2)
ax.add_collection(lc)

ax.set_xlim(x.min(), x.max())
ax.set_ylim(-1.1, 1.1)
ax.set_zlim(0, 2)
plt.show()

不連続変化版

cmapとnorm (BoundaryNorm)を使って配色と閾値を調整します.

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap, BoundaryNorm
from mpl_toolkits.mplot3d.axes3d import Axes3D
from mpl_toolkits.mplot3d.art3d import Line3DCollection

theta  = np.linspace(0, 20 * np.pi, 500)
x = np.cos(theta)
y = np.sin(theta)
z = 0.1 * theta / np.pi
points = np.array([x, y, z]).T.reshape(-1, 1, 3)
segments = np.concatenate([points[:-1], points[1:]], axis=1)

fig = plt.figure()
ax = fig.gca(projection = '3d')
cmap = ListedColormap(['w', 'b', 'r', 'w', 'b'])
norm = BoundaryNorm([1/3, 2/3, 1, 4/3, 5/3, 2], cmap.N)
lc = Line3DCollection(segments, cmap=cmap, norm=norm)
lc.set_array(z)
lc.set_linewidth(2)
ax.add_collection(lc)

ax.set_xlim(x.min(), x.max())
ax.set_ylim(-1.1, 1.1)
ax.set_zlim(0, 2)
plt.show()

参考文献