Repository navigation
Expand file tree
/
Copy pathobjective_function.py
More file actions
135 lines (113 loc) · 4.62 KB
/
Copy pathobjective_function.py
File metadata and controls
135 lines (113 loc) · 4.62 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
from jax import jax, jit
from functools import partial
import jax.numpy as jnp
import numpy as np
"""
This file contains implementation of different problem classes, which are all inherented from the parent class `ObjectiveFunction`. We use JAX to auto-differentiate the functions to get gradient and Hessian.
"""
#======================================== Base & Utils ==========================================#
def generate_symmetric_matrix(smallest_eig, largest_eig, n, rng):
"""
Randomly generate a PSD matrix with given condition number eig_max / eig_min with dimension of n
"""
eigenvalues = np.linspace(smallest_eig, largest_eig, n)
Q = rng.random((n, n))
Q, _ = np.linalg.qr(Q)
A = Q.dot(np.diag(eigenvalues)).dot(Q.T)
return A
class ObjectiveFunction():
"""
This class implements fx, gradient and hessian of fx
It is intentionally set up without internal parameters to keep static for jit
This will precompile gradient and Hessian computation so that they run fast online
"""
@partial(jit, static_argnums=0)
def dfdx(self, x): # function gradient
return jax.grad(self.f)(x)
@partial(jit, static_argnums=0)
def hessian(self, x): # function Hessian
return jax.jacfwd(jax.grad(self.f))(x)
def construct_problem(name, x0):
"""
Initialize a problem object with the user specified problem name and starting point
"""
if name == "P1_quad_10_10":
problem = Problem1()
elif name == "Rosenbrock_2":
problem = Problem7()
elif name == "Rosenbrock_100":
problem = Problem8()
elif name == "Exponential_100":
problem = Problem11()
else: # raise an exception if name is out of the scope
raise Exception("Please pass a proper problem name! Options: ['P1_quad_10_10', 'Rosenbrock_2', 'Rosenbrock_100', 'Exponential_100']")
problem.set_initial_point(x0) # set initial point
return problem
#======================================== Project Problems ==========================================#
# All problem classes have common fields `name` and `x0` indicating the probelm name string and the initial guess
# The field `x0` needs to be set by invoking the `self.set_initial_point` method
class Problem1(ObjectiveFunction):
def __init__(self):
rng = np.random.default_rng(seed = 0) # random number generator
self.q = rng.random(10) # random q
self.Q = generate_symmetric_matrix(1, 10, 10, rng) # random Q with specified condition number
self.dim_n = 10 # dimension
self.name = "P1_quad_10_10"
def set_initial_point(self, x0):
# set up initial point
if not jnp.shape(x0)[0] == self.dim_n: # check dimension
raise Exception("The initial guess x0 should have the correct dimension!")
else:
self.x0 = x0
@partial(jit, static_argnums=0)
def f(self, x): # function value
fx = 0.5 * x.T @ self.Q @ x + self.q.T @ x
return fx
class Problem7(ObjectiveFunction):
def __init__(self):
self.dim_n = 2
self.name = "Rosenbrock_2"
def set_initial_point(self, x0):
# set up initial point
if not jnp.shape(x0)[0] == self.dim_n:
raise Exception("The initial guess x0 should have the correct dimension!")
else:
self.x0 = x0
@partial(jit, static_argnums=0)
def f(self, x):
fx = 0.0
for ii in range(self.dim_n - 1):
fx += 100 * (x[ii+1] - x[ii]**2)**2 + (1 - x[ii])**2
return fx
class Problem8(ObjectiveFunction):
def __init__(self):
self.dim_n = 100
self.name = "Rosenbrock_100"
def set_initial_point(self, x0):
# set up initial point
if not jnp.shape(x0)[0] == self.dim_n:
raise Exception("The initial guess x0 should have the correct dimension!")
else:
self.x0 = x0
@partial(jit, static_argnums=0)
def f(self, x):
fx = 0.0
for ii in range(self.dim_n - 1):
fx += 100 * (x[ii+1] - x[ii]**2)**2 + (1 - x[ii])**2
return fx
class Problem11(ObjectiveFunction):
def __init__(self):
self.dim_n = 100
self.name = "Exponential_100"
def set_initial_point(self, x0):
# set up initial point
if not jnp.shape(x0)[0] == self.dim_n:
raise Exception("The initial guess x0 should have the correct dimension!")
else:
self.x0 = x0
@partial(jit, static_argnums=0)
def f(self, x):
fx = (jnp.exp(x[0]) - 1) / (jnp.exp(x[0]) + 1) + 0.1 * jnp.exp(-x[0])
for ii in np.arange(1, 100):
fx += (x[ii] - 1)**4
return fx