python 3Dグラフ Z方向を値によって色変更したい

Question

地形図のような３Dグラフをプロットした際に、指定したZの値よりも大きいものと小さいもので色分けをしたいと考えています。
もしくは、同じグラフ上にZ=##という平面を描きたいです。

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


# CSVからデータを読み込む
data = pd.read_excel('//tikeizu.xlsx', index_col=0)
print(data.columns.values)
# 3Dグラフの初期化
fig = plt.figure()
ax = fig.gca(projection='3d')

# データの準備
Xgrid = data.columns.values.astype(np.float32)
Ygrid = data.index.values.astype(np.float32)
X, Y = np.meshgrid(Xgrid, Ygrid)
Z = data.values

# プロット
surf = ax.plot_surface(X, Y, Z)


# 表示
plt.show()

調べてもなかなか出てきません。
ご教授お願いいたします。

Answer 1

Zの値よりも大きいものと小さいもので色分けをしたいと考えています。

ある値以上か未満かで色分けする自作のカラーマップを作って、cmap 引数に設定してください。

サンプルコード

z=40 で色分けする例

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


def f(x, y):
    return x ** 2 + y ** 2 + x * y


X, Y = np.mgrid[-6:6:100j, -6:6:100j]
Z = f(X, Y)

fig = plt.figure(figsize=(9, 9), facecolor="w")
ax = fig.add_subplot(111, projection="3d")

### 40 以上は緑、40未満は黄緑で表示する例
thresh = 40  # 色を返る閾値
cmap = colors.ListedColormap(["lightgreen", "green"])
bounds = [Z.min(), thresh, Z.max()]
norm = colors.BoundaryNorm(bounds, cmap.N)

ax.plot_surface(X, Y, Z, cmap=cmap, norm=norm)

plt.show()

Comments

[p_pp]

ありがとうございます！