引数がスカラーである関数の3次元曲面を描画したい

Question

前提・実現したいこと

３次元グラフを出力するプログラムを書いています。２次元メッシュを作成後、関数に代入すると以下のエラーが出ます。関数の引数がスカラーのみしか受け付けないことが原因だと考えられます。関数を書き換えることなく曲面を描画することは可能でしょうか。

発生している問題・エラーメッセージ

TypeError: only size-1 arrays can be converted to Python scalars

該当のソースコード

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

def func(x, y):
    '--省略--'
    return z

x = np.arange(0, 1, 0.1)
y = np.arange(0, 1, 0.1)
X, Y = np.meshgrid(x, y)
Z = func(X, Y) # TypeError

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

ax.plot_wireframe(X, Y, Z)
plt.show()

補足情報（FW/ツールのバージョンなど）

python 3.8.8

Answer 1

funcの詳細不明なので未検証ですがnumpy.vectorizeでいけるかと思います。

vf = np.vectorize(func)
Z = vf(X,Y)

Comments

[cream_puff]

素早い対応ありがとうございました。