def test_numpy(self): import numpy as np from scipy.optimize import rosen, rosen_der from dict_minimize.numpy_api import minimize def rosen_obj(params, shift): val = rosen(params["x_half_a"] - shift) + rosen(params["x_half_b"] - shift) dval = OrderedDict([ ("x_half_a", rosen_der(params["x_half_a"] - shift)), ("x_half_b", rosen_der(params["x_half_b"] - shift)), ]) return val, dval np.random.seed(0) n_a = 3 n_b = 5 shift = -1.0 params = OrderedDict([("x_half_a", np.random.randn(n_a)), ("x_half_b", np.random.randn(n_b))]) params = minimize(rosen_obj, params, args=(shift, ), method="L-BFGS-B", options={"disp": False})
def test_torch(self): import torch from dict_minimize.torch_api import minimize def rosen_obj(params, shift): """Based on augmented Rosenbrock from botorch.""" X, Y = params["x_half_a"], params["x_half_b"] X = X - shift Y = Y - shift obj = 100 * (X[1] - X[0]**2)**2 + 100 * (Y[1] - Y[0]**2)**2 return obj def d_rosen_obj(params, shift): obj = rosen_obj(params, shift=shift) da, db = torch.autograd.grad( obj, [params["x_half_a"], params["x_half_b"]]) d_obj = OrderedDict([("x_half_a", da), ("x_half_b", db)]) return obj, d_obj torch.manual_seed(123) n_a = 2 n_b = 2 shift = -1.0 params = OrderedDict([("x_half_a", torch.randn((n_a, ))), ("x_half_b", torch.randn((n_b, )))]) params = minimize(d_rosen_obj, params, args=(shift, ), method="L-BFGS-B", options={"disp": False})
def test_minimize(x0_dict, args, method, tol): total_el = sum(vv.size for vv, _ in x0_dict.values()) assume(total_el > 0) x0_dict = OrderedDict([(kk, from_np(vv, dd)) for kk, (vv, dd) in x0_dict.items()]) args = tuple(args) def dummy_f(xk, *args_): assert args == args_ validate_solution(x0_dict, xk) # Pass back some arbitrary stuff v = sum(vv.sum() for vv in xk.values()) dv = OrderedDict([kk, vv + 1] for kk, vv in xk.items()) return v, dv def callback(xk): validate_solution(x0_dict, xk) x_sol = minimize(dummy_f, x0_dict, args=args, method=method, tol=tol, callback=callback, options={"maxiter": 10}) validate_solution(x0_dict, x_sol)
def test_minimize_bounded(x0_dict_, args, method, tol): total_el = sum(vv.size for vv, _ in x0_dict_.values()) assume(total_el > 0) args = tuple(args) check_bounds = method in always_respects_bounds x0_dict = OrderedDict() lb_dict = OrderedDict() ub_dict = OrderedDict() for kk, (vv, dd) in x0_dict_.items(): vv = np.sort(vv, axis=0) lb, vv, ub = vv x0_dict[kk] = from_np(vv, dd) lb_dict[kk] = from_np(lb, dd) ub_dict[kk] = from_np(ub, dd) validate_solution(x0_dict, x0_dict, lb_dict, ub_dict) def dummy_f(xk, *args_): assert args == args_ if check_bounds: validate_solution(x0_dict, xk, lb_dict, ub_dict) else: validate_solution(x0_dict, xk) # Pass back some arbitrary stuff v = sum(vv.sum() for vv in xk.values()) dv = OrderedDict([kk, vv + 1] for kk, vv in xk.items()) return v, dv def callback(xk): if check_bounds: validate_solution(x0_dict, xk, lb_dict, ub_dict) else: validate_solution(x0_dict, xk) x_sol = minimize( dummy_f, x0_dict, args=args, lb_dict=lb_dict, ub_dict=ub_dict, method=method, tol=tol, callback=callback, options={"maxiter": 10}, ) if check_bounds: validate_solution(x0_dict, x_sol, lb_dict, ub_dict) else: validate_solution(x0_dict, x_sol)
def test_tensorflow(self): import tensorflow as tf from dict_minimize.tensorflow_api import minimize def rosen_obj(params, shift): """Based on augmented Rosenbrock from botorch.""" X, Y = params["x_half_a"], params["x_half_b"] X = X - shift Y = Y - shift obj = 100 * (X[1] - X[0]**2)**2 + 100 * (Y[1] - Y[0]**2)**2 return obj def d_rosen_obj(params, shift): with tf.GradientTape(persistent=True) as t: t.watch(params["x_half_a"]) t.watch(params["x_half_b"]) obj = rosen_obj(params, shift=shift) da = t.gradient(obj, params["x_half_a"]) db = t.gradient(obj, params["x_half_b"]) d_obj = OrderedDict([("x_half_a", da), ("x_half_b", db)]) del t # Explicitly drop the reference to the tape return obj, d_obj tf.random.set_seed(123) n_a = 2 n_b = 2 shift = -1.0 params = OrderedDict([("x_half_a", tf.random.normal((n_a, ))), ("x_half_b", tf.random.normal((n_b, )))]) params = minimize(d_rosen_obj, params, args=(shift, ), method="L-BFGS-B", options={"disp": False})
def test_jax(self): from jax import random, value_and_grad import jax.numpy as np from dict_minimize.jax_api import minimize def rosen(x): r = np.sum(100.0 * (x[1:] - x[:-1]**2.0)**2.0 + (1 - x[:-1])**2.0, axis=0) return r def rosen_obj(params, shift): val = rosen(params["x_half_a"] - shift) + rosen(params["x_half_b"] - shift) return val n_a = 3 n_b = 5 shift = -1.0 # Jax makes it this simple: d_rosen_obj = value_and_grad(rosen_obj, argnums=0) # Setup randomness in JAX key = random.PRNGKey(0) key, subkey_a = random.split(key) key, subkey_b = random.split(key) params = OrderedDict([ ("x_half_a", random.normal(subkey_a, shape=(n_a, ))), ("x_half_b", random.normal(subkey_b, shape=(n_b, ))) ]) params = minimize(d_rosen_obj, params, args=(shift, ), method="L-BFGS-B", options={"disp": False})