用numpy中的nditer遍历矩阵中的元素

问题由来：

遍历数组时，通常的做法是采用多层for循环嵌套，例如对于一个三维数组，则采用三个for循环进行遍历。但是，这种做法对于高维度矩阵执行力下降。

另外，还有种做法是先把多维度数组展平，操作完后再还原成原来的数组。

这里介绍另一种方法，即用numpy提供的nditer进行迭代。

步骤为：

1、初始化迭代器

it = np.nditer(x, flags=['multi_index'], op_flags=['readwrite'])

以上x是要被遍历的矩阵，flags为指定是否为多维度迭代器，针对数组；op_flags为指定对迭代器进行读写操作。注意这里字符要用中括号括起来。

2、使用迭代器进行循环

while not it.finished:

        idx = it.multi_index

        ...

        it.iternext()

这里it.finished返回一个bool值，即当迭代至最后一个元素时为True,其他情况为false

it.multi_index 返回当前迭代的数组的索引  

it.iternext()为对索引进行更新，类似 idx +=1

以下为一个简单的例子

import numpy as np

if __name__ == '__main__':
    a = np.array([[[1, 2, 3], [4, 5, 6]]])

    it = np.nditer(a, flags=['multi_index'], op_flags=['readwrite'])

    while not it.finished:
        print(it.multi_index)
        it.iternext()

运行结果：

用以上遍历方法编写一个数值方法（中心差分法）求梯度的函数，并用该方法求解y = x^2 + y ^2 + z ^2 的最小值

import numpy as np


def grad_numerical(f, x):
    """
    用数值方法(中心差分法)计算函数f在x处的梯度
    :param f: 函数句柄
    :param x: 自变量
    :return: 梯度grad
    """
    # 创建数组迭代器
    it = np.nditer(x, flags=['multi_index'], op_flags=['readwrite'])
    x = np.array(x, dtype=np.float64)  # 把x转换成numpy数组的float64类型
    # 循环计算每个元素的梯度
    h = 1e-4  # x的微小变量
    grads = np.zeros_like(x)  # 梯度的形状和自变量的形状相同
    while not it.finished:
        idx = it.multi_index
        temp = x[idx]
        # 计算f(x + h)
        x[idx] = temp + h
        f_h1 = f(x[idx])
        # 计算f(x - h)
        x[idx] = temp - h
        f_h2 = f(x[idx])

        # 还原x[idx]
        x[idx] = temp

        # 中心差分法计算梯度 g = (f(x + h) - f(x - h))/2h
        grads[idx] = (f_h1 - f_h2) / (2.0 * h)

        # 迭代索引更新
        it.iternext()

    return grads

if __name__ == '__main__':
    # 用梯度下降法计算f = x^2 + y^2 + z^2 的最小值
    iter_num = 100  # 迭代次数
    lr = 0.1  # 学习率

    f = lambda x: np.sum(x ** 2)

    # 初始值
    x0 = np.array([1., -2., 3.])

    for i in range(iter_num):
        grad = grad_numerical(f, x0)
        x0 -= lr * grad

        if i % 10 == 0:  # 每10次返回一次迭代结果
            y = f(x0)
            print(f'第{i}次迭代函数值为{y}, 对应自变量为{x0}')

当然，也可用最前面提到的第二种方法，即将数据展平成1维后再reshape:

import numpy as np


def grad_numerical1(f, x):
    """
    用数值方法(中心差分法)计算函数f在x处的梯度
    :param f: 函数句柄
    :param x: 自变量
    :return: 梯度grad
    """

    x = np.array(x, dtype=np.float64)  # 把x转换成numpy数组的float64类型
    x_shape = x.shape
    x_len = x.size

    x_iter = x.reshape(-1)
    grads = np.zeros_like(x_iter)  # 初始化梯度
    # 循环计算每个元素的梯度
    h = 1e-4  # x的微小变量
    for idx in range(x_len):
        temp = x_iter[idx]
        # 计算f(x + h)
        x_iter[idx] = temp + h
        f_h1 = f(x_iter[idx])
        # 计算f(x - h)
        x_iter[idx] = temp - h
        f_h2 = f(x_iter[idx])

        # 还原x[idx]
        x_iter[idx] = temp

        # 中心差分法计算梯度g = (f(x + h) - f(x - h)) / 2h
        grads[idx] = (f_h1 - f_h2) / (2.0 * h)

    grads = grads.reshape(x_shape)
    return grads

if __name__ == '__main__':
    # 用梯度下降法计算f = x^2 + y^2 + z^2的最小值
    iter_num = 100  # 迭代次数
    lr = 0.1  # 学习率

    f = lambda x: np.sum(x ** 2)

    # 初始值
    x0 = np.array([1., -2., 3.])

    for i in range(iter_num):
        grad = grad_numerical(f, x0)
        x0 -= lr * grad

        if i % 10 == 0:  # 每10次返回一次迭代结果
            y = f(x0)
            print(f'第{i}次迭代函数值为{y}, 对应自变量为{x0}')