matplotlib.Axes3dで物理っぽい図を作る

Matplotlibの3Dプロットは問題が多くあまり実用的とは思えないのですが，使いたいときもあります．
物理の教科書っぽいきれいな原理図を作れるといいなあと思って試してみました．

環境

matplotlib 3.5.1
python 3.10.4
Ubuntu 20.04 on win

結果の図

作ってみたのは以下のような図．ポイントは，

• デフォルトになっている薄気味悪い灰色の壁を何とかする
• デフォルトの軸を消して新しい矢印をつける
• 3dのラジアルプロットっぽい棒にカラーマップを付ける
  というところ．

コード

t = np.linspace(0, 2*π, 200)
x = np.cos(3*t)
y = np.sin(3*t)
cmap = cm.jet

# 元々ある3dの軸を無理やり消す
plt.rcParams['axes.linewidth'] = 0

# Axes3dのプロパティ
conf = {
    "projection": "3d",
    "proj_type": "persp", # ortho or persp
    "box_aspect": (2.8 ,1.2, 1),
    "elev": 40,
    "azim": -60, # ax.view_init(elev=40, azim=-60) これでもok
    "axisbelow": True,
    "facecolor": "w", 
    "xticks": [],
    "yticks": [],
    "zticks": [],
    }

fig = plt.figure(figsize=(10,10))
ax = fig.add_subplot(111, **conf)

# 元々ある3dの壁を消す
ax.xaxis.pane.fill = False
ax.yaxis.pane.fill = False
ax.zaxis.pane.fill = False

# x, y, z軸の矢印の追加
arrow_conf = {
    "head_width": 0.15,
    "color": "k",
}
xax = mpl.patches.FancyArrow(x=0, y=0, dx=1.5, dy=0, **arrow_conf)
yax = mpl.patches.FancyArrow(x=0, y=0, dx=0, dy=1.5, **arrow_conf)
zax = mpl.patches.FancyArrow(x=0, y=0, dx=7, dy=0  , **arrow_conf)
ax.add_patch(xax)
ax.add_patch(yax)
ax.add_patch(zax)
mpl3d.art3d.patch_2d_to_3d(xax, z=0, zdir="x")
mpl3d.art3d.patch_2d_to_3d(yax, z=0, zdir="x")
mpl3d.art3d.patch_2d_to_3d(zax, z=0, zdir="z")

# データのプロット
ax.plot(t, x, y, "k--")
q = ax.quiver(t, 0, 0, 0, x, y, cmap=cmap, arrow_length_ratio=0, lw=0.5)
q.set_array(t)

留意点

• 灰色の壁の消し方
  ax.xaxis.pane.fill = Falseでできます．元々調べた方法はxaxisではなくw_axisを使うものでしたが，こちらはdeprecatedです．

• set_visibleは使えない
  まず，元々あるx,y,z軸を非表示にする方法が見つからなかった．
  2DだとSpinesやxaxisなどのオブジェクトがset_visible()というメソッドを持っているので，これにFalseを渡すとかがよくある軸の消し方．ところが３Dだとspinesやxaxisにset_visible(False)とすると

UnboundLocalError: local variable 'axis_bb' referenced before assignment

となってまともに動かない．結局，力技でxaxisのlinewidthをゼロにして対処した．これもなぜかrcParams経由でないと反映されなかった．理由はよくわからない．

• FancyArrowの追加
  ax.add_patch(Patch)で矢印をaxに追加する．その後でmpl3d.art3d.patch_2d_to_3d(patch, z, zdir)とするとうまく矢印が追加できる．
  zdirは面の指定（"z"とするとx-y平面が指定される）で，zで面からのオフセットを指定する．
  patch_2d_to_3dはPatchオブジェクトをPatch3Dオブジェクトに変換するものらしいが，変換後のオブジェクトを返すわけではなく，すでにあるインスタンスに直接作用しているようである．詳しくはソースコード参照．
  ちなみにこの方法を見出すまではax.quiver()（ベクトル場の描画などに使われるあれ）を使って矢印を描画するという方法を試していたが，矢印の頭のサイズをどうしてもうまく設定できずに諦めた．

• ステムプロットっぽいやつ
  座標軸から放射状に棒を伸ばすプロット方法はけっこう見かけると思う．何と呼ぶのかわからないが，ステムプロットというものと少し似ている．これはax.quiver()でarrow_length_ratio=0として矢印の頭を消すことで対応できる．ちなみに前述したとおり，矢印の頭のサイズの計算にも問題があるようで，とくに軸同士のスケールが異なる場合はきれいに表示されない．ので，消すしかないと思われる．

• 3D版quiver()のグラデーション
  quiverはもともとcmapによるグラデーション表示に対応しているはずだが，３Dの場合はそうはいかないらしい．これはset_array()というメソッドで対応できる．これはquiver()がもともと持っているはずのCという引数に対応するもので，カラーマップをデータに対してノーマライズするための配列を受け付けるもの（その配列の値で色が決まる）．

結論

matplotlibの３Dプロット用APIを私は応援している．軽いしね．