RosenBrock関数を実装してみた（Python）

非線形関数のRosenBrock関数について実装を行いたかったが意外と勾配も含めて実装されている文献がなかったので実装の備忘録を書きました。

入力の次元がD次元の時、関数値はこの通りですね。
$$ f(x) = \sum_{i=1}^{D-1} [100{(x_{i+1} – {x_i}^2 )}^2 + {(x_i – 1)}^2 ] $$

とりあえず実装してみました。
加えて勾配ですがこちらは入力の次元分だけ出力値が出るはずです。
同じく入力がD次元あったとき微分する次元（\(x_d\)）で関数が異なります。
1)d=0
2)d=D
3)それ以外
のときで場合わけできます。
1)のとき
$$ g(x) = -400(x_{i+1} – {x_i}^2 )x_i + 2(x_i – 1) $$
2)のとき
$$ g(x) = 200(x_{i+1} – {x_i}^2 )x_i $$
3)のとき
$$ g(x) = -400(x_{i+1} – {x_i}^2 )x_i + 2(x_i – 1) + 200(x_i+1 – {x_i}^2 )x_i $$

class RosenBrock(models.Model):
    def __init__(self, name="RosenBrock", noise_value=np.array([0,0]), var=1):
        super(RosenBrock, self).__init__(name=name)
        self.w_star = np.ones(2)
        self.noise_value = noise_value

    def f_opt(self, w):
        w = np.array(w)
        d = w.shape[0]
        self.w_star = np.ones(d)
        for i in range(0, d - 1):
            tmp_1 = 100 * (w[i + 1] - w[i] ** 2) ** 2
            tmp_2 = (w[i] - 1) ** 2
            tmp = tmp_1 + tmp_2

        self.noise_value = self.add_noise(tmp)
        tmp = tmp + self.noise_value

        return tmp

    def g_opt(self, w):
        w = np.array(w)
        d = w.shape[0]
        self.w_star = np.ones(d)
        tmp = np.zeros(w.shape)
        for i in range(d):
            if i == 0:
                tmp[i] = 100 * (-4) * w[i] * (w[i + 1] - w[i] ** 2) + 2 * (w[i] - 1)
            elif i == d - 1:
                tmp[i] = 100 * 2 * (w[i] - w[i - 1] ** 2)
            else:
                tmp[i] = 100 * (-4) * w[i] * (w[i + 1] - w[i] ** 2) + 2 * (w[i] - 1) + 100 * 2 * (w[i] - w[i - 1] ** 2)

        tmp = tmp + self.noise_value

        return tmp

f = model_opt.RosenBrock()
x = np.arange(-0.5, 1.5, 0.1)
y = np.arange(-0.5, 1.5, 0.1)
X, Y = np.meshgrid(x, y)
Z = f.f_opt([X,Y])

fig = plt.figure()
ax = Axes3D(fig)
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("f(x, y)")

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