tryout cost_set_slice

pyextra/acados_template/acados_ocp_solver_pyx.pyx

@@ -278,7 +278,40 @@ cdef class AcadosOcpSolverFast:
                     self.nlp_dims, self.nlp_out, stage, field, <void *> &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, <void *> &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, <void *> &value[0][0])

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):

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))