複素数対応版を探すのに苦労したのでメモ。
精度については調査していない。

複素数の使い方

虚数単位としてjを付けるだけ。.realと.imagで実部、虚部を取り出せる。

1+2j

ガンマ関数

scipy.specialのgammaを使う。ガンマ関数の逆数rgammaも用意されている。
array_like型を引数に取れる。

mpmath.gammaもある。

(リーマン)ゼータ関数

mpmath.zetaを使えば複素数領域の計算が可能。

なお、scipy.specialのzeta, zetac関数は記事作成時点(2023/1/8)では複素数非対応。

インストール

pip install mpmath

プロット

matplotlib.pyplotを使う。

ソースコード

scipyのガンマ関数

表示の都合で値6.0で上限を設定しています。

from mpl_toolkits import mplot3d
import numpy as np
import matplotlib.pyplot as plt
from scipy.special import *

Z_LIMIT = 6.0

def f(x, y):
    return np.clip(abs(gamma(x + y*1j)),0,Z_LIMIT)

x = np.linspace(-4, 4, 160)
y = np.linspace(-2, 2, 80)

X, Y = np.meshgrid(x, y)
Z = f(X, Y)
fig = plt.figure()
ax = plt.axes(projection='3d')
ax.view_init(elev=20, azim=225)
ax.set_zlim3d(top=Z_LIMIT)
ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap='viridis', edgecolor='none')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')

plt.show()

mpmathのガンマ関数

mpmathのgamma関数の引数は、array_like型を取れないので、for文を使って配列要素ごとに計算して設定してます。

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

import mpmath
mpmath.mp.dps = 20

Z_LIMIT = 6.0

def f(x, y):
    z = np.zeros((len(x), len(x[0])))
    for i in range(len(x)):
        for k in range(len(x[0])):
            z[i][k] = abs(mpmath.gamma(x[i][k] + y[i][k]*1j))
            if z[i][k] > Z_LIMIT :
                z[i][k] = Z_LIMIT
    return z

x = np.linspace(-4, 4, 160)
y = np.linspace(-2, 2, 80)

X, Y = np.meshgrid(x, y)
Z = f(X, Y)
fig = plt.figure()
ax = plt.axes(projection='3d')
ax.view_init(elev=20, azim=225)
ax.set_zlim3d(top=Z_LIMIT)
ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap='viridis', edgecolor='none')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')

plt.show()

ガンマ関数（の絶対値）のプロット結果

scipyのgammaとmpmath(dps=20設定時)のgammaの差（厳密にはabsで絶対値取ったあとの差なので、参考程度に。）

Pythonでガンマ関数とゼータ関数をプロットするためのメモ書き（複素数対応版）