diff --git a/nosnoc/__init__.py b/nosnoc/__init__.py index 65d1337..69e8522 100644 --- a/nosnoc/__init__.py +++ b/nosnoc/__init__.py @@ -1,5 +1,5 @@ from .auto_model import NosnocAutoModel -from .solver import NosnocSolver, get_results_from_primal_vector +from .solver import NosnocSolver, get_results_from_primal_vector, construct_problem from .problem import NosnocProblem from .model import NosnocModel from .ocp import NosnocOcp diff --git a/nosnoc/model.py b/nosnoc/model.py index dd8ba82..6c7f4a0 100644 --- a/nosnoc/model.py +++ b/nosnoc/model.py @@ -298,6 +298,8 @@ def preprocess_model(self, opts: NosnocOpts): std_compl_res += ca.transpose(lambda_p[ii]) @ (np.ones(n_c_sys[ii]) - alpha_ii) lambda00_expr = ca.vertcat(lambda00_expr, -ca.fmin(self.c[ii], 0), ca.fmax(self.c[ii], 0)) + else: + raise NotImplementedError() # collect all algebraic equations g_z_all = ca.vertcat(g_switching, g_convex, g_lift, self.g_z) # g_lift_forces diff --git a/nosnoc/nosnoc_types.py b/nosnoc/nosnoc_types.py index 5e9aba5..9b7b4b8 100644 --- a/nosnoc/nosnoc_types.py +++ b/nosnoc/nosnoc_types.py @@ -27,6 +27,7 @@ class MpccMode(Enum): ELASTIC_INEQ = auto() ELASTIC_EQ = auto() ELASTIC_TWO_SIDED = auto() + BOOLEAN = auto() # KANZOW_SCHWARTZ = auto() # NOSNOC: 'scholtes_ineq' (3), 'scholtes_eq' (2) # NOTE: tested in simple_sim_tests diff --git a/nosnoc/problem.py b/nosnoc/problem.py index ac4223b..9cd3d64 100644 --- a/nosnoc/problem.py +++ b/nosnoc/problem.py @@ -140,6 +140,19 @@ def create_complementarity(self, x: List[ca.SX], y: ca.SX, sigma: ca.SX, tau: ca for j in range(n): for x_i in x: g_comp[j + 3 * n] = opts.fb_ip_aug2_weight * (g_comp[j]) * ca.sqrt(1 + (x_i[j] - y[j])**2) + elif opts.mpcc_mode == MpccMode.BOOLEAN: + bigM = 1e5 + n_z = y.numel() + z = ca.SX.sym("z", n_z) + self.add_variable(z, self.ind_bool, np.zeros(n_z), + np.ones(n_z), 0.5 * np.ones(n_z)) + _x = casadi_sum_list([x_i for x_i in x]) + g_comp = ca.vcat( + [bigM * z - y, # y < bigM * z -> 0 < ... + y + bigM * z, + bigM * (1 - z) - _x, + _x + bigM * (1 - z)] + ) n_comp = casadi_length(g_comp) if opts.mpcc_mode in [MpccMode.SCHOLTES_INEQ, MpccMode.ELASTIC_INEQ]: @@ -155,11 +168,12 @@ def create_complementarity(self, x: List[ca.SX], y: ca.SX, sigma: ca.SX, tau: ca elif opts.mpcc_mode == MpccMode.ELASTIC_TWO_SIDED: lb_comp = np.hstack((-ca.inf*np.ones((n,)), np.zeros((n,)))) ub_comp = np.hstack((np.zeros((n,)), ca.inf*np.ones((n,)))) + elif opts.mpcc_mode == MpccMode.BOOLEAN: + lb_comp = np.zeros((n_comp,)) + ub_comp = ca.inf * np.ones((n_comp,)) self.add_constraint(g_comp, lb=lb_comp, ub=ub_comp, index=self.ind_comp) - return - class FiniteElementBase(NosnocFormulationObject): @@ -254,14 +268,18 @@ def __init__(self, self.ind_h = [] self.ind_comp = [] + self.ind_bool = [] # create variables - h = ca.SX.sym(f'h_{ctrl_idx}_{fe_idx}') h_ctrl_stage = opts.terminal_time / opts.N_stages h0 = np.array([h_ctrl_stage / np.array(opts.Nfe_list[ctrl_idx])]) - ubh = (1 + opts.gamma_h) * h0 - lbh = (1 - opts.gamma_h) * h0 - self.add_step_size_variable(h, lbh, ubh, h0) + if opts.use_fesd: + self.h = ca.SX.sym(f'h_{ctrl_idx}_{fe_idx}') + ubh = (1 + opts.gamma_h) * h0 + lbh = (1 - opts.gamma_h) * h0 + self.add_step_size_variable(self.h, lbh, ubh, h0) + else: + self.h = h0 if opts.mpcc_mode in [MpccMode.SCHOLTES_EQ, MpccMode.SCHOLTES_INEQ, MpccMode.ELASTIC_INEQ, MpccMode.ELASTIC_EQ]: lb_dual = 0.0 @@ -373,7 +391,7 @@ def __init__(self, ocp.lbx, ocp.ubx, model.x0, -1) def add_step_size_variable(self, symbolic: ca.SX, lb: float, ub: float, initial: float): - self.ind_h = casadi_length(self.w) + self.ind_h = [casadi_length(self.w)] self.w = ca.vertcat(self.w, symbolic) self.lbw = np.append(self.lbw, lb) @@ -411,9 +429,6 @@ def sum_Lambda(self, sys=slice(None)): Lambdas = self.get_Lambdas_incl_last_prev_fe(sys) return casadi_sum_list(Lambdas) - def h(self) -> List[ca.SX]: - return self.w[self.ind_h] - def X_fe(self) -> List[ca.SX]: """returns list of all x values in finite element""" opts = self.opts @@ -424,7 +439,7 @@ def X_fe(self) -> List[ca.SX]: for j in range(opts.n_s): x_temp = self.prev_fe.w[self.prev_fe.ind_x[-1]] for r in range(opts.n_s): - x_temp += self.h() * opts.A_irk[j, r] * self.w[self.ind_v[r]] + x_temp += self.h * opts.A_irk[j, r] * self.w[self.ind_v[r]] X_fe.append(x_temp) X_fe.append(self.w[self.ind_x[-1]]) elif opts.irk_representation == IrkRepresentation.DIFFERENTIAL_LIFT_X: @@ -447,7 +462,7 @@ def forward_simulation(self, ocp: NosnocOcp, Uk: ca.SX, sot: ca.SX) -> None: for j in range(opts.n_s): x_temp = self.prev_fe.w[self.prev_fe.ind_x[-1]] for r in range(opts.n_s): - x_temp += self.h() * opts.A_irk[j, r] * self.w[self.ind_v[r]] + x_temp += self.h * opts.A_irk[j, r] * self.w[self.ind_v[r]] # lifting constraints self.add_constraint(self.w[self.ind_x[j]] - x_temp) @@ -463,13 +478,13 @@ def forward_simulation(self, ocp: NosnocOcp, Uk: ca.SX, sot: ca.SX) -> None: for r in range(opts.n_s): xj += opts.C_irk[r + 1, j + 1] * X_fe[r] Xk_end += opts.D_irk[j + 1] * X_fe[j] - self.add_constraint(self.h() * fj - xj) - self.cost += opts.B_irk[j + 1] * self.h() * qj + self.add_constraint(self.h * fj - xj) + self.cost += opts.B_irk[j + 1] * self.h * qj elif (opts.irk_representation in [IrkRepresentation.DIFFERENTIAL, IrkRepresentation.DIFFERENTIAL_LIFT_X]): - Xk_end += self.h() * opts.b_irk[j] * self.w[self.ind_v[j]] + Xk_end += self.h * opts.b_irk[j] * self.w[self.ind_v[j]] self.add_constraint(fj - self.w[self.ind_v[j]]) - self.cost += opts.b_irk[j] * self.h() * qj + self.cost += opts.b_irk[j] * self.h * qj # continuity condition: end of fe state - final stage state if (not opts.right_boundary_point_explicit or @@ -557,13 +572,13 @@ def step_equilibration(self, sigma_p: ca.SX, tau: ca.SX, s_elastic: Optional[ca. # only step equilibration mode that does not require previous finite element if opts.step_equilibration == StepEquilibrationMode.HEURISTIC_MEAN: h_fe = opts.terminal_time / (opts.N_stages * opts.Nfe_list[self.ctrl_idx]) - self.cost += opts.rho_h * (self.h() - h_fe)**2 + self.cost += opts.rho_h * (self.h - h_fe)**2 return elif not self.fe_idx > 0: return prev_fe: FiniteElement = self.prev_fe - delta_h_ki = self.h() - prev_fe.h() + delta_h_ki = self.h - prev_fe.h if opts.step_equilibration == StepEquilibrationMode.HEURISTIC_DELTA: self.cost += opts.rho_h * delta_h_ki**2 @@ -612,7 +627,6 @@ def __create_control_stage(self, ctrl_idx, prev_fe): ctrl_idx, fe_idx=ii, prev_fe=prev_fe) - self._add_finite_element(fe, ctrl_idx) control_stage.append(fe) prev_fe = fe return control_stage @@ -643,12 +657,18 @@ def __create_primal_variables(self): self.stages.append(stage) prev_fe = stage[-1] - def _add_finite_element(self, fe: FiniteElement, ctrl_idx: int): + def _collect_finite_elements(self): + """Collect all FE""" + for ctrl_idx, stage in enumerate(self.stages): + for fe in stage: + self.__collect_final_element(fe, ctrl_idx) + + def __collect_final_element(self, fe: FiniteElement, ctrl_idx: int): w_len = casadi_length(self.w) self._add_primal_vector(fe.w, fe.lbw, fe.ubw, fe.w0) # update all indices - self.ind_h.append(fe.ind_h + w_len) + self.ind_h.extend(increment_indices(fe.ind_h, w_len)) self.ind_x[ctrl_idx].append(increment_indices(fe.ind_x, w_len)) self.ind_x_cont[ctrl_idx].append(increment_indices(fe.ind_x[-1], w_len)) self.ind_v[ctrl_idx].append(increment_indices(fe.ind_v, w_len)) @@ -659,6 +679,7 @@ def _add_finite_element(self, fe: FiniteElement, ctrl_idx: int): self.ind_lambda_n[ctrl_idx].append(increment_indices(fe.ind_lambda_n, w_len)) self.ind_lambda_p[ctrl_idx].append(increment_indices(fe.ind_lambda_p, w_len)) self.ind_z[ctrl_idx].append(increment_indices(fe.ind_z, w_len)) + self.ind_bool[ctrl_idx].append(increment_indices(fe.ind_bool, w_len)) # TODO: can we just use add_variable? It is a bit involved, since index vectors here have different format. def _add_primal_vector(self, symbolic: ca.SX, lb: np.array, ub, initial): @@ -717,6 +738,7 @@ def __init__(self, opts: NosnocOpts, model: NosnocModel, ocp: Optional[NosnocOcp self.ind_alpha = create_empty_list_matrix((opts.N_stages,)) self.ind_lambda_n = create_empty_list_matrix((opts.N_stages,)) self.ind_lambda_p = create_empty_list_matrix((opts.N_stages,)) + self.ind_bool = create_empty_list_matrix((opts.N_stages,)) self.ind_z = create_empty_list_matrix((opts.N_stages,)) self.ind_elastic = [] @@ -730,6 +752,7 @@ def __init__(self, opts: NosnocOpts, model: NosnocModel, ocp: Optional[NosnocOcp # setup parameters, lambda00 is added later: sigma_p = ca.SX.sym('sigma_p') # homotopy parameter + tau = ca.SX.sym('tau') # homotopy parameter if opts.mpcc_mode in [MpccMode.ELASTIC_TWO_SIDED, MpccMode.ELASTIC_EQ, MpccMode.ELASTIC_INEQ]: # Elasticity parameter s_elastic = ca.SX.sym('s_elastic') @@ -740,7 +763,6 @@ def __init__(self, opts: NosnocOpts, model: NosnocModel, ocp: Optional[NosnocOcp else: s_elastic = None - tau = ca.SX.sym('tau') # homotopy parameter self.p = ca.vertcat(casadi_vertcat_list(model.p_ctrl_stages), sigma_p, tau) # Generate all the variables we need @@ -776,7 +798,7 @@ def __init__(self, opts: NosnocOpts, model: NosnocModel, ocp: Optional[NosnocOcp self.add_fe_constraints(fe, ctrl_idx) if opts.use_fesd and opts.equidistant_control_grid: - self.add_constraint(sum([fe.h() for fe in stage]) - h_ctrl_stage) + self.add_constraint(sum([fe.h for fe in stage]) - h_ctrl_stage) if opts.time_freezing and opts.equidistant_control_grid: # TODO: make t0 dynamic (since now it needs to be 0!) @@ -798,7 +820,8 @@ def __init__(self, opts: NosnocOpts, model: NosnocModel, ocp: Optional[NosnocOcp # terminal constraint and cost # NOTE: this was evaluated at Xk_end (expression for previous state before) # which should be worse for convergence. - x_terminal = self.w[self.ind_x[-1][-1][-1]] + last_fe = self.stages[-1][-1] + x_terminal = last_fe.w[last_fe.ind_x[-1]] g_terminal = ocp.g_terminal_fun(x_terminal, model.p_ctrl_stages[-1], model.v_global) self.add_constraint(g_terminal) self.cost += ocp.f_q_T_fun(x_terminal, model.p_ctrl_stages[-1], model.v_global) @@ -812,9 +835,12 @@ def __init__(self, opts: NosnocOpts, model: NosnocModel, ocp: Optional[NosnocOcp # Terminal numerical time if opts.N_stages > 1 and opts.use_fesd and not opts.equidistant_control_grid: - all_h = [fe.h() for stage in self.stages for fe in stage] + all_h = [fe.h for stage in self.stages for fe in stage] self.add_constraint(sum(all_h) - opts.terminal_time) + # Collect all w + self._collect_finite_elements() + # CasADi Functions self.cost_fun = ca.Function('cost_fun', [self.w, self.p], [self.cost]) self.comp_res = ca.Function('comp_res', [self.w, self.p], [J_comp]) @@ -869,3 +895,25 @@ def is_sim_problem(self): if not self.ocp_trivial: return False return True + + def dump(self, file): + """Dump the problem to a file.""" + import pickle + with open(file, "wb") as f: + data = { + key: self.__dict__[key] for key in self.__dict__ + if key not in [ + "w", "p", "cost", "g", "model", "ocp", "stages", "fe0" + ] + } + data["f"] = ca.Function("f", [self.w, self.p], [self.cost]) + data["g"] = ca.Function("g", [self.w, self.p], [self.g]) + data["w_shape"] = self.w.shape + data["p_shape"] = self.p.shape + + sigma, tau = 0.0, 0.0 + lambda00 = self.model.compute_lambda00(self.opts) + data["p0"] = np.concatenate( + (self.model.p_val_ctrl_stages.flatten(), + np.array([sigma, tau]), lambda00, self.model.x0)) + f.write(pickle.dumps(data)) diff --git a/nosnoc/solver.py b/nosnoc/solver.py index d8f8884..a2b3628 100644 --- a/nosnoc/solver.py +++ b/nosnoc/solver.py @@ -14,6 +14,17 @@ from nosnoc.utils import flatten_layer, flatten, get_cont_algebraic_indices, flatten_outer_layers, check_ipopt_success +def construct_problem(opts: NosnocOpts, model: NosnocModel, ocp: Optional[NosnocOcp] = None) -> NosnocProblem: + # preprocess inputs + opts.preprocess() + model.preprocess_model(opts) + + if opts.initialization_strategy == InitializationStrategy.RK4_SMOOTHENED: + model.add_smooth_step_representation(smoothing_parameter=opts.smoothing_parameter) + + return NosnocProblem(opts, model, ocp) + + class NosnocSolverBase(ABC): @abstractmethod @@ -21,22 +32,11 @@ def __init__(self, opts: NosnocOpts, model: NosnocModel, ocp: Optional[NosnocOcp] = None) -> None: - # preprocess inputs - opts.preprocess() - model.preprocess_model(opts) - - if opts.initialization_strategy == InitializationStrategy.RK4_SMOOTHENED: - model.add_smooth_step_representation(smoothing_parameter=opts.smoothing_parameter) - # store references self.model = model self.ocp = ocp self.opts = opts - - # create problem - problem = NosnocProblem(opts, model, ocp) - self.problem: NosnocProblem = problem - return + self.problem = construct_problem(opts, model, ocp) def set(self, field: str, value: np.ndarray) -> None: """ @@ -49,7 +49,7 @@ def set(self, field: str, value: np.ndarray) -> None: dims = prob.model.dims if field == 'x0': prob.model.x0 = value - elif field == 'x': # TODO: check other dimensions for useful error message + elif field == 'x': # TODO: check other dimensions for useful error message if value.shape[0] == self.opts.N_stages: # Shape is equal to the number of control stages for i, sub_idx in enumerate(prob.ind_x): @@ -503,8 +503,9 @@ def get_results_from_primal_vector(prob: NosnocProblem, w_opt: np.ndarray) -> di time_steps = w_opt[prob.ind_h] else: t_stages = opts.terminal_time / opts.N_stages + time_steps = [] for Nfe in opts.Nfe_list: - time_steps = Nfe * [t_stages / Nfe] + time_steps += [t_stages / Nfe] * Nfe results["time_steps"] = time_steps # results relevant for OCP: