diff --git a/pyextra/acados_template/acados_ocp_solver_pyx.pyx b/pyextra/acados_template/acados_ocp_solver_pyx.pyx index aa272d6728..3716e3ba38 100644 --- a/pyextra/acados_template/acados_ocp_solver_pyx.pyx +++ b/pyextra/acados_template/acados_ocp_solver_pyx.pyx @@ -278,7 +278,40 @@ cdef class AcadosOcpSolverFast: self.nlp_dims, self.nlp_out, stage, field, &value[0]) - def cost_set(self, int stage, str field_, value_): + def cost_set_slice(self, int start_stage, int end_stage, bytes field, value_, bint const=False): + """ + """ + assert len(value_.shape) >= 2 + + cdef int i, idx + cdef int dims[2] + cdef double[:,:,:] value + + value_shape = value_.shape + if len(value_shape) == 2: + value = value_[:,:,None] + value_shape = (value_shape[0], value_shape[1], 0) + else: + value = value_ + + assert value.strides[1] == 8 # list of column-major matrices; [:,::1,:] + + for i in range(start_stage, end_stage): + idx = 0 if const else i + + acados_solver_common.ocp_nlp_cost_dims_get_from_attr(self.nlp_config, \ + self.nlp_dims, self.nlp_out, i, field, &dims[0]) + + if value_shape[1] != dims[0] or value_shape[2] != dims[1]: + raise Exception('AcadosOcpSolver.cost_set(): mismatching dimension', \ + ' for field "{}" with dimension {} (you have {})'.format( \ + field, tuple(dims), value_shape)) + + acados_solver_common.ocp_nlp_cost_model_set(self.nlp_config, \ + self.nlp_dims, self.nlp_in, i, field, &value[idx][0][0]) + + + def cost_set(self, int stage, bytes field, value_): """ Set numerical data in the cost module of the solver. @@ -286,8 +319,6 @@ cdef class AcadosOcpSolverFast: :param field: string, e.g. 'yref', 'W', 'ext_cost_num_hess' :param value: of appropriate size """ - field = field_.encode('utf-8') - cdef int dims[2] acados_solver_common.ocp_nlp_cost_dims_get_from_attr(self.nlp_config, \ self.nlp_dims, self.nlp_out, stage, field, &dims[0]) @@ -306,7 +337,7 @@ cdef class AcadosOcpSolverFast: if value_shape[0] != dims[0] or value_shape[1] != dims[1]: raise Exception('AcadosOcpSolver.cost_set(): mismatching dimension', \ ' for field "{}" with dimension {} (you have {})'.format( \ - field_, tuple(dims), value_shape)) + field, tuple(dims), value_shape)) acados_solver_common.ocp_nlp_cost_model_set(self.nlp_config, \ self.nlp_dims, self.nlp_in, stage, field, &value[0][0]) diff --git a/selfdrive/controls/lib/lateral_mpc_lib/lat_mpc.py b/selfdrive/controls/lib/lateral_mpc_lib/lat_mpc.py index 6162cf68a5..413b0fdb6e 100755 --- a/selfdrive/controls/lib/lateral_mpc_lib/lat_mpc.py +++ b/selfdrive/controls/lib/lateral_mpc_lib/lat_mpc.py @@ -121,9 +121,8 @@ class LateralMpc(): self.x_sol = np.zeros((N+1, X_DIM)) self.u_sol = np.zeros((N, 1)) self.yref = np.zeros((N+1, 3)) - for i in range(N): - self.solver.cost_set(i, "yref", self.yref[i]) - self.solver.cost_set(N, "yref", self.yref[N][:2]) + self.solver.cost_set_slice(0, N, b"yref", self.yref) + self.solver.cost_set(N, b"yref", self.yref[N][:2]) # Somehow needed for stable init for i in range(N+1): @@ -136,10 +135,9 @@ class LateralMpc(): def set_weights(self, path_weight, heading_weight, steer_rate_weight): W = np.asfortranarray(np.diag([path_weight, heading_weight, steer_rate_weight])) - for i in range(N): - self.solver.cost_set(i, 'W', W) + self.solver.cost_set_slice(0, N, b'W', W[None,:,:], const=True) #TODO hacky weights to keep behavior the same - self.solver.cost_set(N, 'W', (3/20.)*W[:2,:2]) + self.solver.cost_set(N, b'W', (3/20.)*W[:2,:2]) def run(self, x0, v_ego, car_rotation_radius, y_pts, heading_pts): x0_cp = np.copy(x0) @@ -147,9 +145,8 @@ class LateralMpc(): self.solver.constraints_set(0, "ubx", x0_cp) self.yref[:,0] = y_pts self.yref[:,1] = heading_pts*(v_ego+5.0) - for i in range(N): - self.solver.cost_set(i, "yref", self.yref[i]) - self.solver.cost_set(N, "yref", self.yref[N][:2]) + self.solver.cost_set_slice(0, N, b"yref", self.yref) + self.solver.cost_set(N, b"yref", self.yref[N][:2]) self.solution_status = self.solver.solve() for i in range(N+1): diff --git a/selfdrive/controls/lib/longitudinal_mpc_lib/long_mpc.py b/selfdrive/controls/lib/longitudinal_mpc_lib/long_mpc.py index df6d8e0fb6..b5f12849f7 100644 --- a/selfdrive/controls/lib/longitudinal_mpc_lib/long_mpc.py +++ b/selfdrive/controls/lib/longitudinal_mpc_lib/long_mpc.py @@ -200,9 +200,8 @@ class LongitudinalMpc(): self.a_solution = [0.0 for i in range(N+1)] self.j_solution = [0.0 for i in range(N)] self.yref = np.zeros((N+1, COST_DIM)) - for i in range(N): - self.solver.cost_set(i, "yref", self.yref[i]) - self.solver.cost_set(N, "yref", self.yref[N][:COST_E_DIM]) + self.solver.cost_set_slice(0, N, b"yref", self.yref) + self.solver.cost_set(N, b"yref", self.yref[N][:COST_E_DIM]) self.x_sol = np.zeros((N+1, X_DIM)) self.u_sol = np.zeros((N,1)) self.params = np.zeros((N+1,3)) @@ -223,29 +222,25 @@ class LongitudinalMpc(): def set_weights_for_lead_policy(self): W = np.asfortranarray(np.diag([X_EGO_OBSTACLE_COST, X_EGO_COST, V_EGO_COST, A_EGO_COST, J_EGO_COST])) - for i in range(N): - self.solver.cost_set(i, 'W', W) + self.solver.cost_set_slice(0, N, b'W', W[None,:,:], const=True) # Setting the slice without the copy make the array not contiguous, # causing issues with the C interface. - self.solver.cost_set(N, 'W', np.copy(W[:COST_E_DIM, :COST_E_DIM])) + self.solver.cost_set(N, b'W', np.copy(W[:COST_E_DIM, :COST_E_DIM])) # Set L2 slack cost on lower bound constraints Zl = np.array([LIMIT_COST, LIMIT_COST, LIMIT_COST, DANGER_ZONE_COST]) - for i in range(N): - self.solver.cost_set(i, 'Zl', Zl) + self.solver.cost_set_slice(0, N, b'Zl', Zl[None,:], const=True) def set_weights_for_xva_policy(self): W = np.asfortranarray(np.diag([0., 10., 1., 10., 1.])) - for i in range(N): - self.solver.cost_set(i, 'W', W) + self.solver.cost_set_slice(0, N, b'W', W[None,:,:], const=True) # Setting the slice without the copy make the array not contiguous, # causing issues with the C interface. - self.solver.cost_set(N, 'W', np.copy(W[:COST_E_DIM, :COST_E_DIM])) + self.solver.cost_set(N, b'W', np.copy(W[:COST_E_DIM, :COST_E_DIM])) # Set L2 slack cost on lower bound constraints Zl = np.array([LIMIT_COST, LIMIT_COST, LIMIT_COST, 0.0]) - for i in range(N): - self.solver.cost_set(i, 'Zl', Zl) + self.solver.cost_set_slice(0, N, b'Zl', Zl[None,:], const=True) def set_cur_state(self, v, a): if abs(self.x0[1] - v) > 1.: @@ -332,9 +327,8 @@ class LongitudinalMpc(): self.yref[:,1] = x self.yref[:,2] = v self.yref[:,3] = a - for i in range(N): - self.solver.cost_set(i, "yref", self.yref[i]) - self.solver.cost_set(N, "yref", self.yref[N][:COST_E_DIM]) + self.solver.cost_set_slice(0, N, b"yref", self.yref) + self.solver.cost_set(N, b"yref", self.yref[N][:COST_E_DIM]) self.accel_limit_arr[:,0] = -10. self.accel_limit_arr[:,1] = 10. x_obstacle = 1e5*np.ones((N+1))