NumPyで勾配をベクトル場図示したい！

Question

オライリー・ジャパン発行の、「ゼロから作るDeepLearning」を独学で勉強中です。
プログラミングは初めてで、pythonに関しても同書で同時に勉強中です。
色々調べたりしながら、見よう見まねで動かして来ましたが、今回は調べど調べど行き詰ってしまいました。。。
NumPyの配列の扱いなど、基礎的な理解が何か抜けているような気がしてなりません。。。
自分の認識のどこが誤りなのか、ご教授頂ければ幸いです。
```###キーワード： python, NumPy, 二次元勾配, ベクトル場で図示
やりたいこと⇒「数値微分を使って、あるエリア内のf(x0, x1) = x0^2 + x1^2の勾配を計算し、ベクトル場として図示する。」

meshgridで作った、【x0の要素数40個】×【x1の要素数40個】の計1600個の格子点に対して、
forを使って1点ずつ数値微分を計算し、gradという名の配列に格納しようとしました。
その際、要素と配列とが混同された表記に見えてしまって混乱しており、質問させていただきました。
以下に問題となるコードを示します。＃部分が質問したい箇所です。

def numerical_diff(f, x):
    h = 1e-4
    diff = np.zeros_like(x)

    for i in range(x.size):
        tmp_val = x[i]
        x[i] = tmp_val + h              　# ←配列xの中のi番目の要素x[i]
        fxh1 = f(x)　　　　　              # ←配列xをfに代入して得た配列fxh1
        x[i] = tmp_val - h                  # ←配列xの中のi番目の要素x[i]
        fxh2 = f(x)                            # ←配列xをfに代入して得た配列fxh2
        diff[i] = (fxh1 - fxh2) / (2*h)  # ←配列diffのi番目の要素diff[i]に、配列を代入しようとしている？！どういうこと？？
        x[i] = tmp_val
    return diff

### 試したこと

fxh1 = f(x)を、元々はfxh1 = f(x[i])と書いていて、エラーが出ていました（fxh2においても同様）。
色々試して、前者のように書けばうまく実行されることが分かりました。
自分の理解では、1600個の各格子点を1点ずつ微分し、格納を行っていくのであるから、
ここは前者表記の様な配列としてではなく、後者表記の様な要素として取り扱うのだと考えました。
しかしながら、今回は前者表記が正しいようでして、何がダメなのかや、理屈が分からず、困っています。。。
（正しく動作した時の出力画像はページ最下部に添付しています。）

###エラーメッセージ

Traceback (most recent call last):
File "C:\Users\kohta_yasuoka\Documents\Python Scripts\training\44_def-gradient(f, x, t).py", line 43, in <module>
grad = numerical_grad(f, np.array([X, Y]))
File "C:\Users\kohta_yasuoka\Documents\Python Scripts\training\44_def-gradient(f, x, t).py", line 30, in numerical_grad
grad[i] = numerical_diff(f, x)
File "C:\Users\kohta_yasuoka\Documents\Python Scripts\training\44_def-gradient(f, x, t).py", line 19, in numerical_diff
fxh1 = f(x[i]) #右辺、f(x[i])ではなぜだめ？
File "C:\Users\kohta_yasuoka\Documents\Python Scripts\training\44_def-gradient(f, x, t).py", line 9, in f
return np.sum(x**2, axis=1)
File "C:\Users\kohta_yasuoka\AppData\Local\Continuum\anaconda3\lib\site-packages\numpy\core\fromnumeric.py", line 1880, in sum
return sum(axis=axis, dtype=dtype, out=out, **kwargs)
File "C:\Users\kohta_yasuoka\AppData\Local\Continuum\anaconda3\lib\site-packages\numpy\core_methods.py", line 32, in _sum
return umr_sum(a, axis, dtype, out, keepdims)
numpy.core._internal.AxisError: axis 1 is out of bounds for array of dimension 0

### 該当のソースコード

```python
import numpy as np
import matplotlib.pylab as plt


def f(x):
    if x.ndim == 1:
        return np.sum(x**2)
    else:
        return np.sum(x**2, axis=1)


def numerical_diff(f, x):
    h = 1e-4
    diff = np.zeros_like(x)

    for i in range(x.size):
        tmp_val = x[i]
        x[i] = tmp_val + h
        fxh1 = f(x) #右辺、f(x[i])ではなぜだめ？
        x[i] = tmp_val - h
        fxh2 = f(x) #右辺、f(x[i])ではなぜだめ？
        diff[i] = (fxh1 - fxh2) / (2*h)
        x[i] = tmp_val
    return diff


def numerical_grad(f, X):
    grad = np.zeros_like(X)
    for i, x in enumerate(X):
        grad[i] = numerical_diff(f, x)
    return grad




x0 = np.arange(-2, 2, 0.1)
x1 = np.arange(-2, 2, 0.1)
X, Y = np.meshgrid(x0, x1)

X = X.flatten()
Y = Y.flatten()

grad = numerical_grad(f, np.array([X, Y])) 

plt.figure()
plt.quiver(X, Y, -grad[0], -grad[1], angles="xy")
plt.xlim([-2, 2])
plt.ylim([-2, 2])
plt.xlabel("x0")
plt.ylabel("x1")
plt.show()

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

Python 3.6.5 |Anaconda, Inc.| (default, Mar 29 2018, 13:32:41) [MSC v.1900 64 bit (AMD64)] on win32
Type "help", "copyright", "credits" or "license" for more information.

Answer 1

何故駄目かというと、単純に f 関数が配列データ（ndarray）をパラメータとして取り、リターン値として各要素の二乗和を返すように設計されている為かと思います。
その為質問にあるように各格子点毎の数値(X1 or X2 の値)をf関数に渡すとエラーとなります。

各格子点毎の数値(X1 or X2 の値)をf関数に渡したいのであれば f関数を

def f(x):
    return x**2

とすると問題なく動作するのでは無いでしょうか

Comments

[vattela]

magichan様
稚拙な質問にもお答えいただき、本当にありがとうございます！
回答を下さってからすぐに頂いたソースコードに変更し、問題ない動作を確認できました！
昨日の時点でお礼を言いたかったのですが、元々の方のコードの意味を理解できずに我流で通して満足するのは、果たして解決と呼べるのかと、杞憂かもしれませんが、そちらの解決にも１日頭をひねってみたのですが、こちらは現時点では未だ解決できておらず、、、(;^_^A
おそらく引数や返り値のデータの内情や次元などを理解できておらず、つまづいているのかなと思っています。ｘの中身が覗けたらなぁと思うのですが、、、(￣▽￣)
ただ、これ以上お返事を先延ばしにするのはとも思い、そちらに関してはまた別の質問という形で対処しようと思います。
ともあれ、我流の方での正しいソースコードの書き方が分かったこと、着眼すべき点はどこかが分かったことを持ちまして、本件は解決とし、頂いた回答をベストアンサーとさせていただきました！
本当にありがとうございました！