-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathnew_optimal_advertising_matern.py
141 lines (120 loc) · 4.9 KB
/
new_optimal_advertising_matern.py
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
136
137
138
139
140
141
import jax
import jax.numpy as jnp
import numpy as np
import pyomo.environ as pyo
from pyomo.opt import TerminationCondition
import jsonargparse
from jax import config
from kernels import integrated_matern_kernel_matrices
from typing import List, Optional
config.update("jax_enable_x64", True)
def optimal_advertising_matern(
rho_hat: float = 0.11,
c: float = 0.5,
beta: float = 0.05,
kappa: float = 0.5,
x_0: float = 0.4,
nu: float = 0.5,
sigma: float = 1.0,
rho: float = 15,
solver_type: str = "ipopt",
train_T: float = 40.0,
train_points: int = 41,
test_T: float = 50.0,
test_points: int = 100,
benchmark_T: float = 60.0,
benchmark_points: int = 300,
train_points_list: Optional[List[float]] = None,
verbose: bool = False,
):
# if passing in `train_points` then doesn't us a grid. Otherwise, uses linspace
if train_points_list is None:
train_data = jnp.linspace(0, train_T, train_points)
else:
train_data = jnp.array(train_points_list)
test_data = jnp.linspace(0, test_T, test_points)
benchmark_grid = jnp.linspace(0, benchmark_T, benchmark_points)
# Construct kernel matrices
N = len(train_data)
K, K_tilde = integrated_matern_kernel_matrices(
train_data, train_data, nu, sigma, rho
)
K = np.array(K) # pyomo doesn't support jax arrays
K_tilde = np.array(K_tilde)
# Create pyomo model and variables
m = pyo.ConcreteModel()
m.I = range(N)
m.alpha_x = pyo.Var(m.I, within=pyo.Reals, initialize=0.0)
m.alpha_mu = pyo.Var(m.I, within=pyo.Reals, initialize=0.0)
m.alpha_u = pyo.Var(m.I, within=pyo.Reals, initialize=0.0)
m.mu_0 = pyo.Var(within=pyo.NonNegativeReals, initialize=0.0)
m.u_0 = pyo.Var(within=pyo.NonNegativeReals, initialize=0.0)
# Map kernels to variables. Pyomo doesn't support mu_0 + K_tilde @ m.alpha_mu
def mu(m, i):
return m.mu_0 + sum(K_tilde[i, j] * m.alpha_mu[j] for j in m.I)
def x(m, i):
return x_0 + sum(K_tilde[i, j] * m.alpha_x[j] for j in m.I)
def dmu_dt(m, i):
return sum(K[i, j] * m.alpha_mu[j] for j in m.I)
def dx_dt(m, i):
return sum(K[i, j] * m.alpha_x[j] for j in m.I)
def u(m, i):
return m.u_0 + sum(K_tilde[i, j] * m.alpha_u[j] for j in m.I)
# Define constraints and objective for model and solve
@m.Constraint(m.I) # for each index in m.I
def dx_dt_constraint(m, i):
return dx_dt(m, i) == (1 - x(m, i))*u(m, i)- beta*x(m,i)
gamma = (beta + rho_hat) / c
@m.Constraint(m.I) # for each index in m.I
def dmu_dt_constraint(m, i):
return dmu_dt(m, i) == -gamma + (rho_hat + beta)*mu(m, i) + mu(m, i)*u(m, i)
@m.Constraint(m.I) # for each index in m.I
def shadow_price(m, i):
return u(m, i)**((1.0-kappa)/kappa) - kappa*mu(m, i)*(1 - x(m, i)) == 0.0
@m.Objective(sense=pyo.minimize)
def min_norm(m): # alpha @ K @ alpha not supported by pyomo
return sum(K[i, j] * m.alpha_mu[i] * m.alpha_mu[j] for i in m.I for j in m.I)+sum(K[i, j] * m.alpha_x[i] * m.alpha_x[j] for i in m.I for j in m.I)
solver = pyo.SolverFactory(solver_type)
options = {
"tol": 1e-6, # Tighten the tolerance for optimality
"dual_inf_tol": 1e-6, # Tighten the dual infeasibility tolerance
"constr_viol_tol": 1e-6, # Tighten the constraint violation tolerance
"max_iter": 1000, # Adjust the maximum number of iterations if needed
}
results = solver.solve(m, tee=verbose, options=options)
if not results.solver.termination_condition == TerminationCondition.optimal:
print(str(results.solver)) # raise exception?
alpha_mu = jnp.array([pyo.value(m.alpha_mu[i]) for i in m.I])
alpha_x = jnp.array([pyo.value(m.alpha_x[i]) for i in m.I])
alpha_u = jnp.array([pyo.value(m.alpha_u[i]) for i in m.I])
u_0 = pyo.value(m.u_0)
mu_0 = pyo.value(m.mu_0)
# Interpolator using training data
@jax.jit
def kernel_solution(test_data):
# pointwise comparison test_data to train_data
K_test, K_tilde_test = integrated_matern_kernel_matrices(
test_data, train_data, nu, sigma, rho
)
mu_test = mu_0 + K_tilde_test @ alpha_mu
x_test = x_0 + K_tilde_test @ alpha_x
u_test = u_0 + K_tilde_test @ alpha_u
return x_test, mu_test, u_test
# Generate test_data and compare to the benchmark
x_test, mu_test, u_test = kernel_solution(test_data)
print(f"solve_time(s) = {results.solver.Time}")
return {
"t_train": train_data,
"t_test": test_data,
"x_test": x_test,
"mu_test": mu_test,
"u_test": u_test,
"alpha_mu": alpha_mu,
"alpha_x": alpha_x,
"mu_0": mu_0,
"u_0": u_0,
"solve_time": results.solver.Time,
"kernel_solution": kernel_solution, # interpolator
}
if __name__ == "__main__":
jsonargparse.CLI(optimal_advertising_matern)