def Single_Shooting_RK4(param, method='IPOPT'):
if param == 'Social_Distancing':
from Parameters.Parameters_Social_Distancing import x0, u_min, u_max, f, X_plot, Q_plot, u0
elif param == 'Vaccination':
from Parameters.Parameters_Vaccination_Flat import x0, u_min, u_max, f, X_plot, Q_plot, u0
elif param == 'Isolation':
from Parameters.Parameters_Isolation import x0, u_min, u_max, f, X_plot, Q_plot, u0
else:
return
U = MX.sym('U', N)
NV = U.shape[0]
Ng = 0
Q = f(x0, U)[1]
prob = {'f': Q, 'x': U, 'g': []}
CB = Iteration_Callback('Singleshoot_CB', NV, Ng, iter_step)
if method == 'IPOPT':
opts_IPOPT['iteration_callback'] = CB
solver = nlpsol('solver', 'ipopt', prob, opts_IPOPT)
elif method == 'SQP':
opts_SQP['iteration_callback'] = CB
solver = nlpsol('solver', 'sqpmethod', prob, opts_SQP)
sol = solver(x0=[u0] * N, lbx=[u_min] * N, ubx=[u_max] * N, lbg=[], ubg=[])
sim_data = {
'U': CB.x_sols,
'lam_x': CB.lam_x_sols,
'lam_g': CB.lam_g_sols,
'X': [X_plot(x0, u_sol) for u_sol in CB.x_sols],
'Q': [Q_plot(x0, u_sol) for u_sol in CB.x_sols],
't': tgrid,
't_M': tgrid_M,
'N': N,
'M': M,
'f': sol['f'],
'f_sols': CB.f_sols
}
fname = 'Single_Shooting_' + method + '_' + param
with open(parent + '/data/' + fname + '.pck', 'wb') as file:
pck.dump(sim_data, file)
return fname
def Multiple_shooting_PMP(param, is_armaijo=False):
if param == 'Social_Distancing':
from Parameters.Parameters_Social_Distancing import u_min, u_max, L, xdot, x, nx, u, M, h, N, x0, u0, f
elif param == 'Vaccination':
from Parameters.Parameters_Vaccination_Flat import u_min, u_max, L, xdot, x, nx, u, M, h, N, x0, u0, f
elif param == 'Isolation':
from Parameters.Parameters_Isolation import u_min, u_max, L, xdot, x, nx, u, M, h, N, x0, u0, f
lbd = MX.sym('lbd', (3, 1))
H = L + lbd.T @ xdot
s = vertcat(x, lbd)
s_dot = vertcat(xdot, -jacobian(H, x).T)
# Functions for minimizing hamiltonian wrt. u
H_u = Function('H_u', [u, s], [H])
grad_H_u = Function('grad_H_u', [u, s], [jacobian(H, u)])
if is_armaijo:
Gu = lambda u0, s: armaijo_newton(
grad_H_u, Function('Fu', [u], [jacobian(grad_H_u, u)]), u0)
else:
Gu = rootfinder('Gu', 'newton', grad_H_u)
def argmin_u(u0, s):
u_opt = Gu(u0, s).full()[0][0]
if u_opt < u_min:
u_opt = u_min
elif u_opt > u_max:
u_opt = u_max
return u_opt
F = Function('f', [s, u], [s_dot, L])
fsk, Sk_plot, Qk_plot = RK4_M_times_plot(F, M, h, nx=2 * nx)
fs, S_plot, Q_plot = integrator_N_times_plot(fsk,
N,
Sk_plot,
Qk_plot,
nx=2 * nx)
U = MX.sym('U', N)
S = MX.sym('S', (2 * nx, N + 1))
# Constraint setup
# X0 = x0:
g = [S[:nx, 0] - x0]
for i in range(N):
g.append(S[:, i + 1] - fsk(S[:, i], U[i])[0])
# Lambda_f = 0:
g.append(S[nx:, -1])
g = vertcat(*g)
# Initial Conditions
lbd0 = [1, 1, 1]
traj_initial = True
S0 = [x0, lbd0]
X0 = f(x0, u0)[0].full().squeeze()[:-nx]
for i in range(N):
if traj_initial:
S0.append(X0[i:i + nx])
else:
S0.append(x0)
S0.append(lbd0)
S0 = np.concatenate(S0)
trajinit = '_initial'
Sk = S0
Sk_diff_norm = diff_tol + 1
max_iter = 100
S = S.reshape((-1, 1))
Fr = Function('Fr', [S, U], [g])
jac_Fr = Function('Fr', [S, U], [jacobian(g, S)])
S_sols = []
U_sols = []
iter = 0
uvec = np.linspace(0.5, u_max, 100)
while (iter < max_iter) and (Sk_diff_norm > diff_tol):
print('Iteration = {}, Sk_diff = {}'.format(iter, Sk_diff_norm))
iter += 1
# Calculate optimal U:
U = []
for i in range(N):
sk = Sk[2 * nx * i:2 * nx * i + 2 * nx]
U.append(argmin_u(u0, sk))
# print(U)
# Solve newton-iteration:
Sk_old = Sk
S_sol = (newton_rhapson(lambda s: Fr(s, U),
lambda s: jac_Fr(s, U),
Sk,
tol=1e-6))
Sk = S_sol[-1]
# Sk_diff_norm = norm_1(np.divide(Sk-Sk_old, diff_scaler))
Sk_diff_norm = norm_1(Sk - Sk_old)
U_sols.extend([U] * len(S_sol))
S_sols.extend(S_sol)
X_sols = [np.reshape(X0, (nx, -1))]
X_sols_raw = [np.reshape(X0, (nx, -1))]
U_sols = [np.array([u0] * N)]
lbd_sols_raw = [np.array([lbd0] * (N + 1))]
X_sols_raw.append([s.reshape((2 * nx, -1))[:nx, :] for s in S_sols])
lbd_sols_raw.append(
[s.reshape((2 * nx, -1))[nx:, :].full() for s in S_sols])
S_sols_reconstructed = [
S_plot(s[:2 * nx], u) for s, u in zip(S_sols, U_sols)
]
X_sols.append(
[s.reshape((2 * nx, -1))[:nx, :-nx] for s in S_sols_reconstructed])
lbd_sols = [
s.reshape((2 * nx, -1))[nx:, :-nx].full() for s in S_sols_reconstructed
]
Q_sols = [Q_plot(s_sol[:2 * nx], u) for s_sol, u in zip(S_sols, U_sols)]
a = 1
sim_data = {
'U': U_sols,
'lam_g': lbd_sols_raw,
'lam_x': lbd_sols_raw,
'X': X_sols,
'Q': Q_sols,
'X_raw': X_sols_raw,
't_M': tgrid_M,
't': tgrid,
'N': N,
'M': M,
'f': sum(Q_sols[-1].full()),
'f_sols': [sum(q.full()) for q in Q_sols]
}
fname = 'Multiple_Shooting_PMP' + '_' + param + trajinit
with open(parent + '/data/' + fname + '.pck', 'wb') as file:
pck.dump(sim_data, file)
return fname
def Multiple_Shooting_RK4(param, method='IPOPT', traj_initial=True):
if param == 'Social_Distancing':
from Parameters.Parameters_Social_Distancing import x0, u_min, u_max, u0, f, fk, X_plot, Q_plot, N, M, nx, N_pop, tgrid_M, tgrid
elif param == 'Vaccination':
from Parameters.Parameters_Vaccination_Flat import x0, u_min, u_max, u0, f, fk, X_plot, Q_plot, N, M, nx, N_pop, tgrid_M, tgrid
elif param == 'Isolation':
from Parameters.Parameters_Isolation import x0, u_min, u_max, u0, f, fk, X_plot, Q_plot, N, nx, M, N_pop, tgrid_M, tgrid
else:
return
U = MX.sym('U', N)
X = MX.sym('X', (nx, N))
# Constraint setup
Q = 0
g = [X[:, 0] - x0]
Xk = x0
for i in range(N - 1):
Xk, Qk = fk(Xk, U[i])
g.append(X[:, i + 1] - Xk)
Q += Qk
V = vertcat(U, X.reshape((-1, 1)))
g = vertcat(*g)
NV = V.shape[0]
Ng = g.shape[0]
Q = f(x0, U)[1]
prob = {'f': Q, 'x': V, 'g': g}
CB = Iteration_Callback('Singleshoot_CB', NV, Ng, iter_step)
if method == 'IPOPT':
opts_IPOPT['iteration_callback'] = CB
solver = nlpsol('solver', 'ipopt', prob, opts_IPOPT)
elif method == 'SQP':
opts_SQP['iteration_callback'] = CB
solver = nlpsol('solver', 'sqpmethod', prob, opts_SQP)
trajinit=''
if traj_initial:
X0 = f(x0, u0)[0].full().squeeze()[:-nx]
trajinit= '_initial'
else:
X0 = x0 * N
U0 = [u0]*N
V0 = [*U0, *X0]
ubv = [u_max] * N + [N_pop] * (nx * N)
lbv = [u_min] * N + [0] * (nx * N)
sol = solver(x0=V0, lbx=lbv, ubx=ubv, lbg=[0] * Ng, ubg=[0] * Ng)
U_sols = [U0, *[v_sol[:N] for v_sol in CB.x_sols]]
sim_data = {'U': U_sols, 'lam_x': CB.lam_x_sols, 'lam_g': CB.lam_g_sols,
'X': [X_plot(x0, u) for u in U_sols],
'Q': [Q_plot(x0, u) for u in U_sols],
'X_raw': [np.vstack([V_sol[N:][::nx],V_sol[N+1:][::nx],V_sol[N+2:][::nx]]) for V_sol in CB.x_sols],
't_M': tgrid_M, 't': tgrid, 'N': N, 'M': M, 'f': sol['f'], 'f_sols': CB.f_sols}
fname = 'Multiple_Shooting_' + method + '_' + param + trajinit
with open(parent + '/data/' + fname + '.pck', 'wb') as file:
pck.dump(sim_data, file)
return fname
def Primal_Dual_Multiple_Shooting(param, traj_initial=True):
if param == 'Social_Distancing':
from Parameters.Parameters_Social_Distancing import N_pop, f, x0, u_min, u_max, N, fk, nx, Q_plot, X_plot, u0
if param == 'Vaccination':
from Parameters.Parameters_Vaccination_Flat import N_pop, f, x0, u_min, u_max, N, fk, nx, Q_plot, X_plot, u0
if param == 'Isolation':
from Parameters.Parameters_Isolation import N_pop, f, x0, u_min, u_max, N, fk, nx, Q_plot, X_plot, u0
method = 'Primal_Dual_Multiple_Shooting'
U = MX.sym('U', N)
X = MX.sym('X', (nx, N))
#Constraint setup
Q = 0
g = [X[:, 0] - x0]
Xk = x0
for i in range(N - 1):
Xk, Qk = fk(Xk, U[i])
g.append(X[:, i + 1] - Xk)
Q += Qk
V = vertcat(U, X.reshape((-1, 1)))
g = vertcat(*g)
Ng = g.shape[0]
h = vertcat(u_min - U, U - u_max)
Nh = h.shape[0]
grad_Phi = jacobian(Q, V)
#Primal Dual variables
s = MX.sym('s', Nh)
lbd = MX.sym('lbd', Ng)
mu = MX.sym('mu', Nh)
tau = MX.sym('tau')
w = vertcat(V, lbd, mu, s)
#Primal Dual formulation
grad_lag = grad_Phi.T + jacobian(g, V).T @ lbd + jacobian(h, V).T @ mu
r = vertcat(grad_lag, g, h + s, mu * s - tau)
Fr = Function('r', [w, tau], [r])
jac_Fr = Function('jac_Fr', [w, tau], [jacobian(r, w)])
#Initial Values
tau = 1
trajinit = ''
U0 = [u0] * N
if traj_initial:
X0 = f(x0, U0)[0].full().squeeze()[:-nx]
trajinit = '_initial'
else:
X0 = x0 * N
lbd0 = [1] * Ng
mu0 = [1] * Nh
s0 = [tau / m for m in mu0]
w0 = np.concatenate([U0, X0, lbd0, mu0, s0]).T
wk = w0
jac_Fr0 = jac_Fr(w0, tau)
init_rank = la.matrix_rank(jac_Fr0)
assert init_rank == jac_Fr0.shape[
0], 'Initial rank deficiency: %i' % init_rank + ', should be %i' % jac_Fr0.shape[
0]
tol = 1e-3
max_iter = 100
wk_diff = 1000
diff_scaler = [u_max - u_min] * N + [N_pop] * (N) * nx + [1
] * (Ng + 2 * Nh)
wk_diff_list = []
is_armaijo = False
wk_list = [w0]
wk_opt_list = [w0]
def root_fun(wk, tau):
f = lambda w: Fr(w, tau)
jac_f = lambda w: jac_Fr(w, tau)
if is_armaijo:
return armaijo_newton(f,
jac_f,
wk,
tol=tol,
max_iter=max_iter,
verbose=True)
else:
return newton_rhapson(f,
jac_f,
wk,
tol=tol,
max_iter=max_iter,
verbose=True)
iter = 0
wk_diff_norm = 1000
while wk_diff_norm > diff_tol or (tau > tau_tol) or (iter > max_iter):
print('tau = {}, norm_1(delta_w) = {}'.format(tau, wk_diff_norm))
wk_old = wk_opt_list[-1]
wk_list.extend(root_fun(wk_old, tau))
wk_opt_list.append(wk_list[-1])
wk_diff = wk_old - wk_opt_list[-1]
wk_diff_list.append(wk_diff)
wk_diff_norm = norm_1(np.divide(wk_diff, diff_scaler))
tau *= tau_factor
wk_list.append(wk)
iter += 1
# Format solution data
separate_w = Function('sep_w', [w], [U, X, lbd, mu, s])
U_list = []
X_list = []
lbd_list = []
mu_list = []
s_list = []
for w_sol in wk_opt_list:
U_opt, X_opt, lbd_opt, mu_opt, s_opt = separate_w(w_sol)
_ = [
x.append(y.full())
for x, y in zip([U_list, X_list, lbd_list, mu_list, s_list],
[U_opt, X_opt, lbd_opt, mu_opt, s_opt])
]
sim_data = {
'U': U_list,
'lam_g': lbd_list,
'lam_x': mu_list,
'X': [X_plot(x0, u) for u in U_list],
'Q': [Q_plot(x0, u) for u in U_list],
'f_sols': [sum(Q_plot(x0, u).full()) for u in U_list],
'f': sum(Q_plot(x0, U_list[-1]).full()),
'X_raw': X_list,
't_M': tgrid_M,
't': tgrid,
'N': N,
'M': M
}
fname = 'Multiple_Shooting_Primal_Dual_' + param + trajinit
with open(parent + '/data/' + fname + '.pck', 'wb') as file:
pck.dump(sim_data, file)
return fname
def Primal_Dual_Single_Shooting(param):
if param == 'Social_Distancing':
from Parameters.Parameters_Social_Distancing import u_min, u_max, L, xdot, x, nx, u, M, h, N, x0, u0, f, X_plot, Q_plot
elif param == 'Vaccination':
from Parameters.Parameters_Vaccination_Flat import u_min, u_max, L, xdot, x, nx, u, M, h, N, x0, u0, f, X_plot, Q_plot
elif param == 'Isolation':
from Parameters.Parameters_Isolation import u_min, u_max, L, xdot, x, nx, u, M, h, N, x0, u0, f, X_plot, Q_plot
U = MX.sym('U', N)
X, Q = f(x0, U)
h = vertcat(u_min - U, U - u_max)
Nh = h.shape[0]
s = MX.sym('s', Nh)
mu = MX.sym('mu', Nh)
tau = MX.sym('tau')
w = vertcat(U, mu, s)
grad_Phi = jacobian(Q, U)
grad_lag = grad_Phi.T + jacobian(h, U).T @ mu
r = vertcat(grad_lag, h + s, mu * s - tau)
Fr = Function('r', [w, tau], [r])
jac_Fr = Function('jac_Fr', [w, tau], [jacobian(r, w)])
tau = 1
U0 = [u0] * N
mu0 = [1] * Nh
s0 = [tau / m for m in mu0]
w0 = np.concatenate([U0, mu0, s0]).T
G = rootfinder('root_r', 'newton', Fr)
wk = w0
tol = .5
max_iter = 100
wk_diff = 1000
wk_diff_list = []
wk_list = [w0]
wk_opt = [w0]
while (wk_diff > diff_tol) or (tau > tau_tol):
print('tau = {}, norm_1(delta_w) = {}'.format(tau, wk_diff))
wk_old = wk_opt[-1]
f = lambda x: Fr(x, tau)
grad_f = lambda x: jac_Fr(x, tau)
wk_list.extend(
newton_rhapson(f, grad_f, wk_old, tol=tol, max_iter=max_iter))
wk_opt.append(wk_list[-1].full())
wk_diff = norm_1(wk_old - wk_list[-1])
wk_diff_list.append(wk_diff)
tau *= tau_factor
wk_list.append(wk)
U_sols = [np.array(w[:N]) for w in wk_opt]
X_sols = [X_plot(x0, u).full() for u in U_sols]
Q_sols = [Q_plot(x0, u).full() for u in U_sols]
mu_sols = [w[N:N + Nh] for w in wk_opt]
s_sols = [w[N + Nh:N + 2 * Nh] for w in wk_opt]
sim_data = {
'U': U_sols,
'lam_x': mu_sols,
's': s_sols,
'X': X_sols,
'Q': Q_sols,
'f_sols': [sum(q) for q in Q_sols],
'f': sum(Q_sols[-1]),
't_M': tgrid_M,
'N': N,
'M': M
}
fname = 'Single_Shooting_Primal_Dual_' + param
with open(parent + '/data/' + fname + '.pck', 'wb') as file:
pck.dump(sim_data, file)
return fname