diff --git a/simulations/validate_rdp_gradient_and_hessian.py b/simulations/validate_rdp_gradient_and_hessian.py new file mode 100644 index 0000000..1a57ed3 --- /dev/null +++ b/simulations/validate_rdp_gradient_and_hessian.py @@ -0,0 +1,129 @@ +"""validate gradient and diagonal Hessian of RDP""" + +# TODO: add symmetric weights + +import cupy as xp + + +def get_d_s(x): + """get backward and forward differences and sums for each dimension of an array x + using "edge" padding + """ + x_padded = xp.pad(x, 1, mode="edge") + + d = xp.zeros((2 * x.ndim,) + x.shape, dtype=x.dtype) + s = xp.zeros((2 * x.ndim,) + x.shape, dtype=x.dtype) + + for i in range(x.ndim): + # diff / sum with "backward" neighbor + sl = x.ndim * [slice(1, -1)] + sl[i] = slice(0, -2) + sl = tuple(sl) + + d[2 * i, ...] = x - x_padded[sl] + s[2 * i, ...] = x + x_padded[sl] + + # diff / sum with "forward" neighbor + sl = x.ndim * [slice(1, -1)] + sl[i] = slice(2, None) + sl = tuple(sl) + + d[2 * i + 1, ...] = x - x_padded[sl] + s[2 * i + 1, ...] = x + x_padded[sl] + + return d, s + + +def rdp(x, gamma=2.0, eps=1e-1): + d, s = get_d_s(x) + phi = s + gamma * xp.abs(d) + eps + + tmp = (d**2) / phi + + return float(tmp.sum()) + + +def rdp_grad(x, gamma=2.0, eps=1e-1): + d, s = get_d_s(x) + phi = s + gamma * xp.abs(d) + eps + + tmp = d * (2 * phi - (d + gamma * xp.abs(d))) / (phi**2) + + return 2 * tmp.sum(axis=0) + + +# def rdp_potential(xj, xk, gamma=2.0, eps=1e-1): +# d = xj - xk +# s = xj + xk +# phi = s + gamma * xp.abs(d) + eps +# +# c = (d**2) / phi +# g = 2 * d * (2 * phi - (d + gamma * xp.abs(d))) / (phi**2) +# +# return c, g +# +# +# def rdp2d(x, gamma=2.0, eps=1e-1): +# +# n0, n1 = x.shape +# +# val = 0.0 +# grad = xp.zeros_like(x) +# +# for j in range(n0): +# for k in range(n1): +# jm = j - 1 +# jp = j + 1 +# km = k - 1 +# kp = k + 1 +# +# x0 = x[j, k] +# +# if jm >= 0: +# c, g = rdp_potential(x0, x[jm, k], gamma=gamma, eps=eps) +# val += c +# grad[j, k] += g +# if jp < n0: +# c, g = rdp_potential(x0, x[jp, k], gamma=gamma, eps=eps) +# val += c +# grad[j, k] += g +# if km >= 0: +# c, g = rdp_potential(x0, x[j, km], gamma=gamma, eps=eps) +# val += c +# grad[j, k] += g +# if kp < n1: +# c, g = rdp_potential(x0, x[j, kp], gamma=gamma, eps=eps) +# val += c +# grad[j, k] += g +# +# return val, grad + + +if __name__ == "__main__": + xx = xp.arange(2 * 3 * 4, dtype=xp.float64).reshape(2, 3, 4) + 1 + + gamma = 2.0 + eps = 0.1 + + h = 1e-7 + + f = rdp(xx, gamma=gamma, eps=eps) + g = rdp_grad(xx, gamma=gamma, eps=eps) + + g_num = xp.zeros_like(xx) + + for i in range(xx.shape[0]): + for j in range(xx.shape[1]): + for k in range(xx.shape[2]): + xxp = xx.copy() + xxp[i, j, k] += h + + xxm = xx.copy() + xxm[i, j, k] -= h + + fp = rdp(xxp, gamma=gamma, eps=eps) + fm = rdp(xxm, gamma=gamma, eps=eps) + + g_num[i, j, k] = (fp - fm) / (2 * h) + + assert xp.all(xp.isclose(g, g_num))