listによる、3次元プロットについての問題点

Question

結果から言うと、
ある場所に依存した確率を自分で算出することができました。
ただ、3次元プロットをしようとして、躓いています。
コード自体は、正直汚いと思われ、これから改良していく予定です。

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

#環境設定
n=2  #tの範囲
m=4  #xの範囲

P = [[-1/2, 1/2, 1/2, 1/2],[0,0,0,0],[0,0,0,0],[0,0,0,0]]
Q = [[0,0,0,0],[1/2, -1/2, 1/2, 1/2],[0,0,0,0],[0,0,0,0]]
R = [[0,0,0,0],[0,0,0,0],[1/2, 1/2, -1/2, 1/2],[0,0,0,0]]
S = [[0,0,0,0],[0,0,0,0],[0, 0, 0, 0],[1/2, 1/2, 1/2, -1/2]]

x_list=[]#xline
y_list=[]#yline
t_list=[]#time
p_list=[]#probability
s_list=[]#state
a = 1#1/math.sqrt(2)
b = 0#1j/math.sqrt(2)
c = 0
d = 0

for i in range(0,2*m+1):
    if i == m:
        phi = [a ,b, c, d]
    else:
        phi = [0, 0, 0, 0]
    p = np.dot(phi,np.conj(phi))

    x_list.append(i) #xの座標
    y_list.append(i) #yの座標

    s_list.append(phi)
    p_list.append(p)

for t in range(0, n+1):
    t_list.append(t)
    if t == 0:
        s_list
        p_list
    else:
        next_s_list = [0]*len(s_list)
        for x in range(0,2*m+1):
            if x == 0:
                for y in range(0,2*m+1):   #x=0, 0<= y =<2*m
                    if y == 0:
                        next_s_list[y] = np.inner(P, s_list[x+1]) + np.inner(R, s_list[y+1])
                    elif y == 2*m:
                        next_s_list[y] = np.inner(P, s_list[x+1]) + np.inner(S, s_list[y-1])
                    else:
                        next_s_list[y] = np.inner(P, s_list[x+1]) + np.inner(R,s_list[y+1]) +np.inner(S, s_list[y-1])
            elif x == 2*m:                  #x=2*m, 0<=y<=2*m
                for y in range(0, 2*m+1):
                    if y == 0:
                        next_s_list[y] = np.inner(Q, s_list[x-1]) + np.inner(R, s_list[y+1])
                    elif y == 2*m:
                        next_s_list[y] = np.inner(Q, s_list[x-1]) + np.inner(S, s_list[y-1])
                    else:
                        next_s_list[y]  = np.inner(Q, s_list[x-1]) + np.inner(R,s_list[y+1]) +np.inner(S, s_list[y-1])
            else:
                if y == 0:
                    next_s_list[y] = np.inner(Q,s_list[x-1]) + np.inner(P,s_list[x+1])
                elif y == 2*m:
                    next_s_list[y] = np.inner(Q,s_list[x-1]) + np.inner(P,s_list[x+1]) + np.inner(S,s_list[y-1])
                else:
                    next_s_list[y]  = np.inner(P, s_list[x+1]) + np.inner(Q, s_list[x-1]) + np.inner(R, s_list[y+1]) + np.inner(S,s_list[y-1])
            p_list[x] = np.dot(next_s_list[y], np.conj(next_s_list[y]))
        s_list = next_s_list

    print(t,p_list)


#3Dplot

fig = plt.figure()
ax = Axes3D(fig)

ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("probability")

ax.set_xlim(3*m,-m)
ax.set_ylim(-m,3*m)
ax.set_zlim(0,1)
ax.plot(x_list, y_list, p_list, color ="red", linewidth=1)
plt.show()

今回はコード自体はあまり関係ないと思われますが、一応掲載しときます。確率は

t=0 x_list=y_list=[0,1,2,3,4,5,6,7,8] p_list=[0, 0, 0, 0, 1, 0, 0, 0, 0]
t=1 x_list=y_list=[0,1,2,3,4,5,6,7,8] p_list=[0.0, 0.0, 0.0, 0.25, 0.0, 0.25, 0.0, 0.0, 0.0]
t=2 x_list=y_list=[0,1,2,3,4,5,6,7,8] p_list=[0.0, 0.0, 0.0625, 0.0, 0.125, 0.0, 0.0625, 0.0, 0.0]

となり、計算結果は間違えていません。
このとき、x_list=y_list=[0, 1, 2, 3, 4, 5, 6, 7, 8]に対応しており、
t=1(空白は確率０)のとき
|y/x|0|1|2|3|4|5|6|7|8|
|--|--:|
|0
|1
|2
|3|||||1/4|
|4||||1/4||1/4
|5|||||1/4|
|6
|7
|8

のように、x_list,y_listの中の座標を確率と対応させると、確かに上図のようなプロットが出来るはずなのですが、
x_list,y_listの中の順番とp_listの順番が1:1対応してしまい以下のような図になってしまいました。
原因はわかっているのですが、どのように修正すれば上図のようなプロットに到達できるのか、アドバイス等お願いします。

すこしわかりにくい説明になってしまいました。
状況が理解できない場合、説明します。
よろしくお願い致します。

Answer 1

計算部に関してはあまり理解はしていないのですが・・。
質問の表のようなデータをプロットしたいのであれば、１次元として保持している確立データ(p_list)を２次元データとして保つ必用があるのではないでしょうか。

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

m=4

x_list = range(0, 2*m+1)
y_list = range(0, 2*m+1)
p_list = np.array(
    [[0,0,0,0,0,0,0,0,0],
     [0,0,0,0,0,0,0,0,0],
     [0,0,0,0,0,0,0,0,0],
     [0,0,0,0,0.25,0,0,0,0],
     [0,0,0,0.25,0,0.25,0,0,0],
     [0,0,0,0,0.25,0,0,0,0],
     [0,0,0,0,0,0,0,0,0],
     [0,0,0,0,0,0,0,0,0],
     [0,0,0,0,0,0,0,0,0]]
)

X,Y = np.meshgrid(x_list,y_list)

fig = plt.figure()
ax = Axes3D(fig)

ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("probability")

ax.set_xlim(2*m,0)
ax.set_ylim(0,2*m)
ax.set_zlim(0,1)
ax.plot_wireframe(X, Y, p_list)
plt.savefig('out.png')
plt.show()

【結果】

Comments

[KSwordOfHaste]

plot_surfaceって回答してしまったですが、ワイヤーフレームならplot_wireframeですね！

[magichan]

確かにそうですね。
ただ個人的には、surface と wireframe は好みの問題かと思いますので、surfaceでも全く問題ないかと思います。

[Fallout_18]

んーと、上記のようにp_listにひとつひとつプロットしていくということですか？

[Fallout_18]

新たな2次元配列をもってきて、時間ごとに更新しながら代入していくやり方の方がよさそうですね。

[Fallout_18]

magichanさん、上のp_listは1次元から2次元配列を自分で手作業で代入して変えるしか、やり方がありませんでしょうか？

[Fallout_18]

調べたら、できるみたいなのでちょっとやってみます

[magichan]

返信おそくなりました。すみません
ワイヤーフレームのような３次元のグラフを書くためには、各x,yに対するP値のマトリックスを全て埋める必要があります。
で、その方法ですが、一次元であたえられたＰ値のリストが各x,y にどのように対応しているのかを私が理解しておりませんので、申し訳ありませんが具体的にに回答することができません。とりあえずｘとｙを引数にとりP値を返す関数 f(x,y) 等を定義することで、簡単に記述できるかと思います。

[Fallout_18]

お返事ありがとうございます。
おっしゃる通り、pをx,yの関数で定義して行うことが、解決策になりますね。やってみます。

Answer 2

自分もmatplotlibはサンプル真似るレベルですので間違ってるかも知れませんが...

そもそもplotの1回の呼び出しでは1本の折れ線グラフしか描けないのではないでしょうか。これを用いて確率をX-Y平面上に充填したようなグラフを描画するにはX or Y軸にそって複数本の走査線のようなグラフを別個に描画する必要があるような気がします。

つまり以下のようなイメージです。

...
ax.plot([2, 3, 4, 5, 6], [3, 3, 3, 3, 3], [0, 0,   1/4, 0, 0])
ax.plot([2, 3, 4, 5, 6], [4, 4, 4, 4, 4], [0, 1/4, 0,   1/4, 0])
ax.plot([2, 3, 4, 5, 6], [5, 5, 5, 5, 5], [0, 0,   1/4, 0, 0])
...
plt.show()

もし2次元の格子のグラフをお望みならplot_surfaceを用いるのだと思います。

自分がイメージしたのは
https://matplotlib.org/mpl_toolkits/mplot3d/tutorial.html
にあるサンプルのうち以下のグラフです。
https://matplotlib.org/_images/wire3d_demo1.png

サンプルコードを見ますとplot_surfaceの引数のx, y, zはそれぞれXY平面の各格子上の頂点の座標を2次元に並べた形で与えることがわかります。

もし確率がx, yについての関数f(x, y)として定義できるなら簡単で、

(A) x, yについてはnumpy.meshgridを用いて展開
(B) ff = np.vectorize(f)でスカラー引数前提の関数fを配列前提のものに置き換える
(C) ffに(A)で生成したx, yを与えてz = ff(x, y)とする
(D) ax.plot_surface(x, y, z, ...)

とすれば描画できます。(A)～(C)は計算上の都合でこのとおりにやりにくいなら、同じ結果となるよううに独自に計算してももちろんかまわないです。

Answer 3

出力する配列の次元が違っていたので、そこを修正して解決できました。