３Dグラフのアニメーションを作成する（Python matplotlibの使い方）

f:id:Start_python:20200115144233g:plain

はじめに

前回に続き、今度はmatplotlibモジュールを使って三次元グラフを作ってみます。
「from mpl_toolkits.mplot3d import Axes3D」を使用します。
サイコロの出た目を、前回の目をｘ軸、前々回の目をｙ軸にしたとき、ｚ軸に偶数なら青色で奇数なら赤色で出た目の合計回数グラフで表します。

プログラムのコード

import random
import numpy as np
import matplotlib.pylab as plt
import matplotlib.animation as animation
from mpl_toolkits.mplot3d import Axes3D

# Qテーブルの初期化
q_table_1 = np.zeros((6, 6))
q_table_2 = np.zeros((6, 6))

# ３Dグラフの準備
fig = plt.figure()

x = [1, 2, 3, 4, 5, 6]
y = [1, 2, 3, 4, 5, 6]
X, Y = np.meshgrid(x, y)

a1 = random.randint(0, 5)
a0 = random.randint(0, 5)

def update(i):
    global a0, a1, a2

    plt.cla()

    a2 = a1
    a1 = a0
    a0 = random.randint(0, 5)

    if a0 == 1 or a0 == 3 or a0 == 5:
        q_table_1[a1, a2] += 1
    if a0 == 0 or a0 == 2 or a0 == 4:
        q_table_2[a1, a2] += 1

    Z1 = q_table_1
    Z2 = q_table_2

    ax = Axes3D(fig)
    
    plt.title('i=' + str(i) + '   [ ' + str(a0+1) + ' ]')

    ax.plot_wireframe(X, Y, Z1, color='b')
    ax.plot_wireframe(X, Y, Z2, color='r')

ani = animation.FuncAnimation(fig, update, interval = 100)
plt.show()