def _test_helper(self, f, x, u, df_dx, df_du, debug=False):
if debug:
breakpoint()
f_xu = f(x, u)
input_x = x
input_u = u
with auto_diff.AutoDiff(x, u) as (x, u):
y, (J_fx, J_fu) = auto_diff.get_value_and_jacobians(f(x, u))
self._assertAllClose(y, f_xu)
self._assertAllClose(J_fx, df_dx)
self._assertAllClose(J_fu, df_du)
u = input_u
with auto_diff.AutoDiff(input_x) as x:
y, J_fx = auto_diff.get_value_and_jacobian(f(x, u))
self._assertAllClose(y, f_xu)
self._assertAllClose(J_fx, df_dx)
x = input_x
with auto_diff.AutoDiff(input_u) as u:
y, J_fu = auto_diff.get_value_and_jacobian(f(x, u))
self._assertAllClose(y, f_xu)
self._assertAllClose(J_fu, df_du)
# Some bugs only appeared with rectangular Jacobians.
A_x = np.random.rand(input_x.shape[0], 3 * input_x.shape[0])
b_x = np.random.rand(input_x.shape[0], 1)
affine_x = np.linalg.lstsq(A_x, input_x - b_x, rcond=None)[0]
A_u = np.random.rand(input_u.shape[0], 3 * input_x.shape[0])
b_u = np.random.rand(input_u.shape[0], 1)
affine_u = np.linalg.lstsq(A_u, input_u - b_u, rcond=None)[0]
df_dx = df_dx @ A_x
df_du = df_du @ A_u
with auto_diff.AutoDiff(affine_x, affine_u) as (x, u):
y, (J_fx, J_fu) = auto_diff.get_value_and_jacobians(
f(A_x @ x + b_x, A_u @ u + b_u))
self._assertAllClose(y, f_xu)
self._assertAllClose(J_fx, df_dx)
self._assertAllClose(J_fu, df_du)
with auto_diff.AutoDiff(affine_x) as x:
y, J_fx = auto_diff.get_value_and_jacobian(
f(A_x @ x + b_x, A_u @ affine_u + b_u))
self._assertAllClose(y, f_xu)
self._assertAllClose(J_fx, df_dx)
with auto_diff.AutoDiff(affine_u) as u:
y, J_fu = auto_diff.get_value_and_jacobian(
f(A_x @ affine_x + b_x, A_u @ u + b_u))
self._assertAllClose(y, f_xu)
self._assertAllClose(J_fu, df_du)
def _test_out(self, f, x, df_dx, debug=False):
if debug:
breakpoint()
input_x = x
f_x = f(input_x)
with auto_diff.SparseAutoDiff(input_x) as x:
out_dest = np.empty(f_x.shape)
f(x, out=out_dest)
y, Jf = auto_diff.get_value_and_jacobian(out_dest)
f(x, out=out_dest)
y2, Jf2 = auto_diff.get_value_and_jacobian(out_dest)
self._assertAllClose(f_x, y)
self._assertAllClose(Jf, df_dx)
self._assertAllClose(f_x, y2)
self._assertAllClose(Jf2, df_dx)
def _test_helper(self, f, x, df_dx, debug=False):
if debug:
breakpoint()
input_x = x
f_x = f(input_x)
with auto_diff.AutoDiff(input_x) as x:
y, Jf = auto_diff.get_value_and_jacobian(f(x))
self._assertAllClose(y, f_x)
self._assertAllClose(Jf, df_dx)
# Some bugs only appeared with rectangular Jacobians.
A = np.random.rand(input_x.shape[0], 3 * input_x.shape[0])
b = np.random.rand(input_x.shape[0], 1)
x = np.linalg.lstsq(A, input_x - b, rcond=None)[0]
df_dx = df_dx @ A
with auto_diff.AutoDiff(x) as x:
y, Jf = auto_diff.get_value_and_jacobian(f(A @ x + b))
self._assertAllClose(y, f_x)
self._assertAllClose(Jf, df_dx)
import auto_diff
import numpy as np
def fs(x):
# From http://fourier.eng.hmc.edu/e176/lectures/NM/node21.html Example 1
# f0 = 3*x_0 - cos(x_1*x_2)--1.5
# f1 = 4*x_0^2-625*x_1^2+2*x-1
# f2 = 20*x_2+exp(-x_0*x_1)+9
f0 = 3 * x[0] - np.cos(x[1] * x[2]) - 3 / 2 # Equation 1
f1 = 4 * x[0]**2 - 625 * x[1]**2 + 2 * x[2] - 1 # Equation 2
f2 = 20 * x[2] + np.exp(-x[0] * x[1]) + 9 # Equation 3
res = [f0, f1, f2] # Create Array of Results
return res # Return Jax Array
x = [0, 0, 0]
with auto_diff.AutoDiff(x) as x:
f_eval = fs(x)
y, Jf = auto_diff.get_value_and_jacobian(f_eval)
# print(auto_diff.jacobian(fs))