Pythonで球面調和関数をピラミッド型にプロットする話

コードの解説は後ほど書きますが先にコードだけ公開します。
このソースコードを利用するときは一言連絡していただけると幸いです。

遊び方

• lmaxの値を自由にいじってあそぶ。大きい値を選びすぎるとなんかうまく出力されませんがlmax=10ぐらいだといい感じになります。
• plt.subplots_adjust(wspace=0, hspace=-0.95)はいい感じに、hspaceの値を設定してあげてください。

from scipy.special import sph_harm
import numpy as np

import matplotlib
import matplotlib.pyplot as plt
import matplotlib.cm as cm

def Y_lm(l, m, θ, φ):
    return sph_harm(m, l, φ, θ)

def real_Ylm(l, m, theta, phi):
    
    if m == 0:
        Y = Y_lm(l, m, theta, phi)
    elif m > 0:
        # m>0の場合、以下の線形結合を取る
        Y = 1/np.sqrt(2)*(Y_lm(l,-m, theta, phi) + (-1)**m * Y_lm(l, m, theta, phi))
    else:
        # m<0のとき、下の解に対応させる
        Y = 1.j/np.sqrt(2)*(Y_lm(l ,m, theta, phi) - (-1)**(-m) * Y_lm(l, -m, theta, phi))
    
    return np.real(Y)

# 極座標(r,Θ,φ)からデカルト座標(x,y,z)への射影
def Polar_to_Descartes(r, theta, phi):
    x = r*np.sin(theta)*np.cos(phi)
    y = r*np.sin(theta)*np.sin(phi)
    z = r*np.cos(theta)
    return x, y, z

def my_plot_surface(l, m, fig, ax, 
                    F, θ, φ,
                    c_F, c_type='bwr', c_alpha=1,             # 色設定
                    ax_lim=0.5, ax_title=None,               # ax設定
                    cb_shrink=0.5, cb_pad=0.1, cb_name=None): # color bar設定


    # 色の正規化
    norm = matplotlib.colors.Normalize(vmin=np.min(c_F), vmax=np.max(c_F))  # 最大値を1に規格化
    
    # グラデーションの設定
    cmap = cm.get_cmap(c_type)
    
    # サーフェスプロット
    x, y, z = Polar_to_Descartes(F, θ, φ)  # Fが正の実数でなければエラー
    
    if l == 0 and m == 0:
        ax.plot_surface(x,y,z, color='red', alpha=c_alpha, linewidth=0)
    else:
        ax.plot_surface(x, y, z, 
                    facecolors=cmap(norm(c_F)),  # c_Fが実数でなければエラー
                    alpha=c_alpha, linewidth=0)
        
    # 軸範囲、ラベル、グラフタイトル
    ax.set_xlim(-ax_lim, ax_lim)
    ax.set_ylim(-ax_lim, ax_lim)
    ax.set_zlim(-ax_lim, ax_lim)
    ax.set_title(ax_title)

# 座標
theta = np.linspace(0, np.pi, 100)
phi = np.linspace(0, 2*np.pi, 200)
θ, φ = np.meshgrid(theta, phi)  # 2Dデータに変換

lmax = 10
J = []
J.append(lmax+1)
for k in range(1,lmax+1,1):
    for kk in range((lmax+1)+(2*lmax+1)*k-k,(lmax+1)+(2*lmax+1)*k+k+1,1):
        J.append(kk)
i=0
fig = plt.figure(figsize=(100,100))

for l in range(lmax+1):
    for m in range(-l,l+1,1):
        ax = fig.add_subplot(lmax+1,2*lmax+1,J[i], projection='3d')
        F = np.abs(real_Ylm(l, m, θ, φ))**2
        c_F = real_Ylm(l, m, θ, φ)
        my_plot_surface(l,m,fig, ax, F, θ, φ, c_F)
        ax.tick_params(labelbottom=False,
            labelleft=False,
            labelright=False,
            labeltop=False)
        ax.tick_params(bottom=False,
            left=False,
            right=False,
            top=False)
        plt.subplots_adjust(wspace=0, hspace=-0.95)
        i = i+1