クリックした座標近傍の複素関数の値をプロットする - Python matplotlib

の続き

キャプチャ

操作方法

• 画面左端のグラフ上で座標を左クリックで選択
• 右クリックまたはウィンドウを閉じる操作で終了
• 無操作タイムアウトでの終了は86400秒(丸一日)

左端のグラフから順に

• $s$
• $\xi(s) =\pi^{-\frac{s}{2}}\Gamma\left(\frac{s}{2}\right)\zeta(s)$
• $|\xi(s))|$ (横軸は$Im(s)$)
• $|\xi(s))|$ (横軸は$Re(s)$)

ソースコード

import matplotlib.pyplot as plt
from matplotlib import _pylab_helpers
# import cmath
import numpy as np
import mpmath
mpmath.mp.dps = 12

def zeta_relist_im_abs(re_list,im):
    z_re = np.zeros(len(re_list))
    z_im = np.zeros(len(re_list))
    z_abs = np.zeros(len(re_list))
    for i in range(len(re_list)):
        s = re_list[i] + im*1j
        tmp = mpmath.zeta(s)*mpmath.gamma(0.5*s)*mpmath.power(mpmath.pi,-0.5*s)
        z_re[i] = tmp.real
        z_im[i] = tmp.imag
        z_abs[i] = abs(tmp)
        # z_arg[i] = cmath.phase(tmp)
    return z_re,z_im,z_abs

def zeta_re_imlist_abs(re,im_list):
    z_re = np.zeros(len(im_list))
    z_im = np.zeros(len(im_list))
    z_abs = np.zeros(len(im_list))
    for i in range(len(im_list)):
        s = re + im_list[i]*1j
        tmp = mpmath.zeta(s)*mpmath.gamma(0.5*s)*mpmath.power(mpmath.pi,-0.5*s)
        z_re[i] = tmp.real
        z_im[i] = tmp.imag
        z_abs[i] = abs(tmp)
        # z_arg[i] = cmath.phase(tmp)
    return z_re,z_im,z_abs


def main():
    RE_N=11
    IM_N=51
    # RE_AMP=0.1
    IM_AMP=1
    
    fig,axes = plt.subplots(1, 4, figsize=(13,5), gridspec_kw= {'width_ratios': [1, 3, 2, 2]})
    axes[0].set_xticks([0.5,0.75,1])
    axes[0].set_xlim((0.5,1))
    axes[0].set_yticks([0,14,14.5,20])
    axes[0].set_ylim((0,20))
    axes[0].grid()
    axes[1].grid()
    axes[2].grid()
    axes[3].grid()

    # ハンドル取得のため一度plot
    lines00, = axes[0].plot([0,1], [0,0],"r") # s  Re part sweep
    lines01, = axes[0].plot([0,1], [0,0],"b") # s  Im part sweep s_re
    lines03, = axes[0].plot([0,1], [0,0],"g") # s  Im part sweep s_re max
    lines10, = axes[1].plot([0,1], [0,0],"r") # zeta_xi(s) Re part sweep
    lines11, = axes[1].plot([0,1], [0,0],"b") # zeta_xi(s) Im part sweep
    lines13, = axes[1].plot([0,1], [0,0],"g") # zeta_xi(s) Im part sweep s_re max
    lines21, = axes[2].plot([0,1], [0,0],"b") # abs(zeta(s)) Im part sweep
    lines23, = axes[2].plot([0,1], [0,0],"g") # abs(zeta(s)) Im part sweep s_re max
    # lines24, = axes[2].plot([0,1], [0,0],"y") # diff of abs(zeta(s))
    lines3,  = axes[3].plot([0,1], [0,0],"r") # abs(zeta(s)) Re part sweep 

    while True:
        manager = _pylab_helpers.Gcf.get_active()
        if manager is None:
            print('closed')
            break

        points = plt.ginput(n=1, timeout=86400, mouse_add=1, mouse_pop=2, mouse_stop=3)
        # a = plt.ginput(n=-1, mouse_add=1, mouse_pop=3, mouse_stop=2)
        # n=-1でインプットが終わるまで座標を取得
        # mouse_addで座標を取得（左クリック）
        # mouse_popでUndo（ミドルクリック）
        # mouse_stopでインプットを終了する（右クリック）

        if len(points)==1:
            re = 0.5
            re_max,im = points[0]
            print(re_max, im)
            
            re_list = np.zeros(RE_N)
            im_list_dummy = np.zeros(RE_N)

            re_list_dummy  = np.zeros(IM_N)
            re_list_dummy3 = np.zeros(IM_N)
            im_list = np.zeros(IM_N)

            for i in range(RE_N):
                re_list[i] = re + (re_max-re)*i/(RE_N-1)
                im_list_dummy[i] = im

            for i in range(IM_N):
                re_list_dummy[i]  = re
                re_list_dummy3[i] = re_list[-1]
                im_list[i] = im + IM_AMP*(i-(IM_N-1)/2)/((IM_N-1)/2)

            axes[0].set_ylim(None)
            lines00.set_data(re_list       , im_list_dummy)
            lines01.set_data(re_list_dummy , im_list      )
            lines03.set_data(re_list_dummy3, im_list      )
            
            # Re sweep
            z_re1,z_im1,z_absr = zeta_relist_im_abs(re_list, im)

            # Im sweep
            z_re2,z_im2,z_abs2 = zeta_re_imlist_abs(re, im_list)
            z_re4,z_im4,z_abs4 = zeta_re_imlist_abs(re_list[-1], im_list) # [-1] means last element of the list

            lines10.set_data(z_re1, z_im1)
            lines11.set_data(z_re2, z_im2)
            lines13.set_data(z_re4, z_im4)

            # z_abs_diff = z_abs4-z_abs2
            lines21.set_data(im_list, z_abs2)
            lines23.set_data(im_list, z_abs4)
            #lines24.set_data(im_list, z_abs_diff)
            
            #tmp_abs_max1 = max(abs(z_abs_diff.min()),abs(z_abs_diff.max()))
            tmp_abs_max2 = abs(z_abs2.max())
            tmp_abs_max1 = abs(z_abs4.max())
            tmp_abs_max  = max(tmp_abs_max1,tmp_abs_max2)
            #tmp_min      = min(0,z_abs_diff.min())
            axes[2].set_xlim(im_list[0],im_list[-1])
            #axes[2].set_ylim(tmp_min,tmp_abs_max)
            axes[2].set_ylim(0,tmp_abs_max)

            lines3.set_data(re_list, z_absr)
            axes[3].set_xlim(re_list[0],re_list[-1])
            tmp_abs_max = abs(z_absr.max())
            tmp_abs_min = abs(z_absr.min())
            axes[3].set_ylim(tmp_abs_min,tmp_abs_max)

            tmp_abs_max1 = max(abs(z_re1.min()),abs(z_re1.max()),abs(z_re2.min()),abs(z_re2.max()),abs(z_re4.min()),abs(z_re4.max()))
            tmp_abs_max2 = max(abs(z_im1.min()),abs(z_im1.max()),abs(z_im2.min()),abs(z_im2.max()),abs(z_im4.min()),abs(z_im4.max()))
            
            axes[1].set_xlim((-tmp_abs_max1, tmp_abs_max1))
            axes[1].set_ylim((-tmp_abs_max2, tmp_abs_max2))
            #tmp_abs_max  = max(tmp_abs_max1,tmp_abs_max2)
            #axes[1].set_xlim((-tmp_abs_max, tmp_abs_max))
            #axes[1].set_ylim((-tmp_abs_max, tmp_abs_max))
        else:
            print('quit by right click')
            break

        # plt.savefig('fig_test.png')
        # plt.show()
        # 次の描画(≒操作受付)までの待ち時間(秒)
        plt.pause(0.05)


if __name__ == '__main__':
    main()

参考記事

matplotlibで座標を取得