def mx_to_cx(name: str, function: Union[Callable, SX, MX], *all_param: Any) -> Function: """ Add to the pool of declared casadi function. If the function already exists, it is skipped Parameters ---------- name: str The unique name of the function to add to the casadi functions pool function: Callable The biorbd function to add all_param: dict Any parameters to pass to the biorbd function """ from ..optimization.optimization_variable import OptimizationVariable, OptimizationVariableList from ..optimization.parameters import Parameter, ParameterList cx_types = OptimizationVariable, OptimizationVariableList, Parameter, ParameterList mx = [ var.mx if isinstance(var, cx_types) else var for var in all_param ] cx = [ var.mapping.to_second.map(var.cx) for var in all_param if isinstance(var, cx_types) ] return biorbd.to_casadi_func(name, function, *mx)(*cx)
def plot_com(x, nlp): com_func = biorbd.to_casadi_func("CoMPlot", nlp.model.CoM, nlp.states["q"].mx, expand=False) com_dot_func = biorbd.to_casadi_func("Compute_CoM", nlp.model.CoMdot, nlp.states["q"].mx, nlp.states["qdot"].mx, expand=False) q = nlp.states["q"].mapping.to_second.map(x[nlp.states["q"].index, :]) qdot = nlp.states["qdot"].mapping.to_second.map( x[nlp.states["qdot"].index, :]) return np.concatenate( (np.array(com_func(q)[2, :]), np.array(com_dot_func(q, qdot)[2, :])))
def main(): """ Generate random data, then create a tracking problem, and finally solve it and plot some relevant information """ # Define the problem biorbd_model = biorbd.Model("arm26.bioMod") final_time = 0.5 n_shooting_points = 30 use_residual_torque = True # Generate random data to fit t, markers_ref, x_ref, muscle_excitations_ref = generate_data( biorbd_model, final_time, n_shooting_points) # Track these data biorbd_model = biorbd.Model( "arm26.bioMod" ) # To allow for non free variable, the model must be reloaded ocp = prepare_ocp( biorbd_model, final_time, n_shooting_points, markers_ref, muscle_excitations_ref, x_ref[:biorbd_model.nbQ(), :], use_residual_torque=use_residual_torque, kin_data_to_track="q", ) # --- Solve the program --- # sol = ocp.solve(show_online_optim=True) # --- Show the results --- # q = sol.states["q"] n_q = ocp.nlp[0].model.nbQ() n_mark = ocp.nlp[0].model.nbMarkers() n_frames = q.shape[1] markers = np.ndarray((3, n_mark, q.shape[1])) symbolic_states = MX.sym("x", n_q, 1) markers_func = biorbd.to_casadi_func("ForwardKin", biorbd_model.markers, symbolic_states) for i in range(n_frames): markers[:, :, i] = markers_func(q[:, i]) plt.figure("Markers") n_steps_ode = ocp.nlp[0].ode_solver.steps + 1 if ocp.nlp[ 0].ode_solver.is_direct_collocation else 1 for i in range(markers.shape[1]): plt.plot(np.linspace(0, 2, n_shooting_points + 1), markers_ref[:, i, :].T, "k") plt.plot(np.linspace(0, 2, n_shooting_points * n_steps_ode + 1), markers[:, i, :].T, "r--") plt.xlabel("Time") plt.ylabel("Markers Position") # --- Plot --- # plt.show()
def states_to_markers(biorbd_model, states): nq = biorbd_model.nbQ() n_mark = biorbd_model.nbMarkers() q = cas.MX.sym("q", nq) markers_func = biorbd.to_casadi_func("makers", biorbd_model.markers, q) return np.array(markers_func(states[:nq, :])).reshape((3, n_mark, -1), order="F")
def generate_data(biorbd_model, tf, x0, t_max, n_shoot, noise_std, show_plots=False): def pendulum_ode(t, x, u): return np.concatenate( (x[nq:, np.newaxis], qddot_func(x[:nq], x[nq:], u)))[:, 0] nq = biorbd_model.nbQ() q = cas.MX.sym("q", nq) qdot = cas.MX.sym("qdot", nq) tau = cas.MX.sym("tau", nq) qddot_func = biorbd.to_casadi_func("forw_dyn", biorbd_model.ForwardDynamics, q, qdot, tau) # Simulated data dt = tf / n_shoot controls = np.zeros( (biorbd_model.nbGeneralizedTorque(), n_shoot)) # Control trajectory controls[0, :] = (-np.ones(n_shoot) + np.sin(np.linspace(0, tf, num=n_shoot))) * t_max states = np.zeros((biorbd_model.nbQ() + biorbd_model.nbQdot(), n_shoot)) # State trajectory for n in range(n_shoot): sol = solve_ivp(pendulum_ode, [0, dt], x0, args=(controls[:, n], )) states[:, n] = x0 x0 = sol["y"][:, -1] states[:, -1] = x0 markers = states_to_markers(biorbd_model, states[:biorbd_model.nbQ(), :]) # Simulated noise np.random.seed(42) noise = (np.random.randn(3, biorbd_model.nbMarkers(), n_shoot) - 0.5) * noise_std markers_noised = markers + noise if show_plots: q_plot = plt.plot(states[:nq, :].T) dq_plot = plt.plot(states[nq:, :].T, "--") name_dof = [name.to_string() for name in biorbd_model.nameDof()] plt.legend(q_plot + dq_plot, name_dof + ["d" + name for name in name_dof]) plt.title("Real position and velocity trajectories") plt.figure() marker_plot = plt.plot(markers[1, :, :].T, markers[2, :, :].T) plt.legend(marker_plot, [i.to_string() for i in biorbd_model.markerNames()]) plt.gca().set_prop_cycle(None) plt.plot(markers_noised[1, :, :].T, markers_noised[2, :, :].T, "x") plt.title("2D plot of markers trajectories + noise") plt.show() return states, markers, markers_noised, controls
def impact(transition, all_pn): """ A discontinuous function that simulates an inelastic impact of a new contact point Parameters ---------- transition: PhaseTransition A reference to the phase transition all_pn: PenaltyNodeList The penalty node elements Returns ------- The difference between the last and first node after applying the impulse equations """ ocp = all_pn[0].ocp if ocp.nlp[transition.phase_pre_idx].states.shape != ocp.nlp[transition.phase_post_idx].states.shape: raise RuntimeError( "Impact transition without same nx is not possible, please provide a custom phase transition" ) # Aliases nlp_pre, nlp_post = all_pn[0].nlp, all_pn[1].nlp # A new model is loaded here so we can use pre Qdot with post model, this is a hack and should be dealt # a better way (e.g. create a supplementary variable in v that link the pre and post phase with a # constraint. The transition would therefore apply to node_0 and node_1 (with an augmented ns) model = biorbd.Model(nlp_post.model.path().absolutePath().to_string()) if nlp_post.model.nbContacts() == 0: warn("The chosen model does not have any contact") q_pre = nlp_pre.states["q"].mx qdot_pre = nlp_pre.states["qdot"].mx qdot_impact = model.ComputeConstraintImpulsesDirect(q_pre, qdot_pre).to_mx() val = [] cx_end = [] cx = [] for key in nlp_pre.states: cx_end = vertcat(cx_end, nlp_pre.states[key].mapping.to_second.map(nlp_pre.states[key].cx_end)) cx = vertcat(cx, nlp_post.states[key].mapping.to_second.map(nlp_post.states[key].cx)) post_mx = nlp_post.states[key].mx continuity = nlp_post.states["qdot"].mapping.to_first.map( qdot_impact - post_mx if key == "qdot" else nlp_pre.states[key].mx - post_mx ) val = vertcat(val, continuity) name = f"PHASE_TRANSITION_{nlp_pre.phase_idx}_{nlp_post.phase_idx}" func = biorbd.to_casadi_func(name, val, nlp_pre.states.mx, nlp_post.states.mx)(cx_end, cx) return func
def com_velocity_equality(multinode_constraint, all_pn): """ The centers of mass velocity are equals for the specified phases and the specified nodes Parameters ---------- multinode_constraint : MultinodeConstraint A reference to the phase transition all_pn: PenaltyNodeList The penalty node elements Returns ------- The difference between the state after and before """ nlp_pre, nlp_post = all_pn[0].nlp, all_pn[1].nlp states_pre = multinode_constraint.states_mapping.to_second.map( nlp_pre.states.cx_end) states_post = multinode_constraint.states_mapping.to_first.map( nlp_post.states.cx) states_post_sym_list = [ MX.sym(f"{key}", *nlp_post.states[key].mx.shape) for key in nlp_post.states.keys() ] states_post_sym = vertcat(*states_post_sym_list) if states_pre.shape != states_post.shape: raise RuntimeError( f"Continuity can't be established since the number of x to be matched is {states_pre.shape} in the " f"pre-transition phase and {states_post.shape} post-transition phase. Please use a custom " f"transition or supply states_mapping") pre_com_dot = nlp_pre.model.CoMdot( states_pre[nlp_pre.states["q"].index, :], states_pre[nlp_pre.states["qdot"].index, :]).to_mx() post_com_dot = nlp_post.model.CoMdot( states_post_sym_list[0], states_post_sym_list[1]).to_mx() pre_states_cx = nlp_pre.states.cx_end post_states_cx = nlp_post.states.cx return biorbd.to_casadi_func( "com_dot_equality", pre_com_dot - post_com_dot, states_pre, states_post_sym, )(pre_states_cx, post_states_cx)
def get_penalty_value(ocp, penalty, t, x, u, p): val = penalty.type.value[0](penalty, PenaltyNodeList(ocp, ocp.nlp[0], t, x, u, []), **penalty.params) if isinstance(val, float): return val states = ocp.nlp[0].states.cx if ocp.nlp[0].states.cx.shape != ( 0, 0) else ocp.cx(0, 0) controls = ocp.nlp[0].controls.cx if ocp.nlp[0].controls.cx.shape != ( 0, 0) else ocp.cx(0, 0) parameters = ocp.nlp[0].parameters.cx if ocp.nlp[ 0].parameters.cx.shape != (0, 0) else ocp.cx(0, 0) return biorbd.to_casadi_func("penalty", val, states, controls, parameters)(x[0], u[0], p)
def torque_bounds(x, min_or_max, nlp, minimal_tau=None): func = biorbd.to_casadi_func("torqueMaxPlot", nlp.model.torqueMax, nlp.states["q"].mx, nlp.states["qdot"].mx, expand=False) q = nlp.states["q"].mapping.to_second.map(x[nlp.states["q"].index, :]) qdot = nlp.states["qdot"].mapping.to_second.map( x[nlp.states["qdot"].index, :]) dof = [3, 5, 6, 7] res = np.array( func(q, qdot)[dof, ::2] if min_or_max == 0 else func(q, qdot)[dof, 1::2]) if minimal_tau is not None: res[res < minimal_tau] = minimal_tau return res
def add_casadi_func(self, name: str, function: Union[Callable, SX, MX], *all_param: Any) -> casadi.Function: """ Add to the pool of declared casadi function. If the function already exists, it is skipped Parameters ---------- name: str The unique name of the function to add to the casadi functions pool function: Callable The biorbd function to add all_param: dict Any parameters to pass to the biorbd function """ if name in self.casadi_func: return self.casadi_func[name] else: mx = [var.mx if isinstance(var, OptimizationVariable) else var for var in all_param] self.casadi_func[name] = biorbd.to_casadi_func(name, function, *mx) return self.casadi_func[name]
def _set_penalty_function(self, ocp, objective, fcn: Union[MX, SX], expand: bool = False): # Do not use nlp.add_casadi_func because all functions must be registered state_cx = ocp.cx(0, 0) control_cx = ocp.cx(0, 0) param_cx = ocp.v.parameters_in_list.cx objective.function = biorbd.to_casadi_func(f"{self.name}", fcn[objective.rows, objective.cols], state_cx, control_cx, param_cx, expand=expand) modified_fcn = objective.function(state_cx, control_cx, param_cx) dt_cx = ocp.cx.sym("dt", 1, 1) weight_cx = ocp.cx.sym("weight", 1, 1) target_cx = ocp.cx.sym("target", modified_fcn.shape) modified_fcn = modified_fcn - target_cx modified_fcn = modified_fcn**2 if objective.quadratic else modified_fcn objective.weighted_function = Function( # Do not use nlp.add_casadi_func because all of them must be registered f"{self.name}", [state_cx, control_cx, param_cx, weight_cx, target_cx, dt_cx], [weight_cx * modified_fcn * dt_cx], ) if expand: objective.function.expand() objective.weighted_function.expand()
def generate_data(biorbd_model: biorbd.Model, final_time: float, n_shooting: int, use_residual_torque: bool = True) -> tuple: """ Generate random data. If np.random.seed is defined before, it will always return the same results Parameters ---------- biorbd_model: biorbd.Model The loaded biorbd model final_time: float The time at final node n_shooting: int The number of shooting points use_residual_torque: bool If residual torque are present or not in the dynamics Returns ------- The time, marker, states and controls of the program. The ocp will try to track these """ # Aliases n_q = biorbd_model.nbQ() n_qdot = biorbd_model.nbQdot() n_tau = biorbd_model.nbGeneralizedTorque() n_mus = biorbd_model.nbMuscleTotal() dt = final_time / n_shooting # Casadi related stuff symbolic_q = MX.sym("q", n_q, 1) symbolic_qdot = MX.sym("qdot", n_qdot, 1) symbolic_mus_states = MX.sym("mus", n_mus, 1) symbolic_tau = MX.sym("tau", n_tau, 1) symbolic_mus_controls = MX.sym("mus", n_mus, 1) symbolic_states = vertcat(*(symbolic_q, symbolic_qdot, symbolic_mus_states)) symbolic_controls = vertcat(*(symbolic_tau, symbolic_mus_controls)) symbolic_parameters = MX.sym("u", 0, 0) nlp = NonLinearProgram() nlp.model = biorbd_model nlp.variable_mappings = { "q": BiMapping(range(n_q), range(n_q)), "qdot": BiMapping(range(n_qdot), range(n_qdot)), "tau": BiMapping(range(n_tau), range(n_tau)), "muscles": BiMapping(range(n_mus), range(n_mus)), } markers_func = biorbd.to_casadi_func("ForwardKin", biorbd_model.markers, symbolic_q) nlp.states.append("q", [symbolic_q, symbolic_q], symbolic_q, nlp.variable_mappings["q"]) nlp.states.append("qdot", [symbolic_qdot, symbolic_qdot], symbolic_qdot, nlp.variable_mappings["qdot"]) nlp.states.append("muscles", [symbolic_mus_states, symbolic_mus_states], symbolic_mus_states, nlp.variable_mappings["muscles"]) nlp.controls.append("tau", [symbolic_tau, symbolic_tau], symbolic_tau, nlp.variable_mappings["tau"]) nlp.controls.append( "muscles", [symbolic_mus_controls, symbolic_mus_controls], symbolic_mus_controls, nlp.variable_mappings["muscles"], ) dynamics_func = biorbd.to_casadi_func( "ForwardDyn", DynamicsFunctions.muscles_driven, symbolic_states, symbolic_controls, symbolic_parameters, nlp, False, ) def dyn_interface(t, x, u): u = np.concatenate([np.zeros(n_tau), u]) return np.array(dynamics_func(x, u, np.empty((0, 0)))).squeeze() # Generate some muscle excitations U = np.random.rand(n_shooting, n_mus).T # Integrate and collect the position of the markers accordingly X = np.ndarray((n_q + n_qdot + n_mus, n_shooting + 1)) markers = np.ndarray((3, biorbd_model.nbMarkers(), n_shooting + 1)) def add_to_data(i, q): X[:, i] = q markers[:, :, i] = markers_func(q[:n_q]) x_init = np.array([0] * n_q + [0] * n_qdot + [0.5] * n_mus) add_to_data(0, x_init) for i, u in enumerate(U.T): sol = solve_ivp(dyn_interface, (0, dt), x_init, method="RK45", args=(u, )) x_init = sol["y"][:, -1] add_to_data(i + 1, x_init) time_interp = np.linspace(0, final_time, n_shooting + 1) return time_interp, markers, X, U
def _set_penalty_function(self, all_pn: Union[PenaltyNodeList, list, tuple], fcn: Union[MX, SX]): """ Finalize the preparation of the penalty (setting function and weighted_function) Parameters ---------- all_pn: PenaltyNodeList The nodes fcn: Union[MX, SX] The value of the penalty function """ # Sanity checks if self.transition and self.explicit_derivative: raise ValueError("transition and explicit_derivative cannot be true simultaneously") if self.transition and self.derivative: raise ValueError("transition and derivative cannot be true simultaneously") if self.derivative and self.explicit_derivative: raise ValueError("derivative and explicit_derivative cannot be true simultaneously") if self.transition: ocp = all_pn[0].ocp nlp = all_pn[0].nlp nlp_post = all_pn[1].nlp name = self.name.replace("->", "_").replace(" ", "_") states_pre = nlp.states.cx_end states_post = nlp_post.states.cx controls_pre = nlp.controls.cx_end controls_post = nlp_post.controls.cx state_cx = vertcat(states_pre, states_post) control_cx = vertcat(controls_pre, controls_post) else: ocp = all_pn.ocp nlp = all_pn.nlp name = self.name if self.integrate: state_cx = horzcat(*([all_pn.nlp.states.cx] + all_pn.nlp.states.cx_intermediates_list)) control_cx = all_pn.nlp.controls.cx else: state_cx = all_pn.nlp.states.cx control_cx = all_pn.nlp.controls.cx if self.explicit_derivative: if self.derivative: raise RuntimeError("derivative and explicit_derivative cannot be simultaneously true") state_cx = horzcat(state_cx, all_pn.nlp.states.cx_end) control_cx = horzcat(control_cx, all_pn.nlp.controls.cx_end) param_cx = nlp.cx(nlp.parameters.cx) # Do not use nlp.add_casadi_func because all functions must be registered self.function = biorbd.to_casadi_func( name, fcn[self.rows, self.cols], state_cx, control_cx, param_cx, expand=self.expand ) if self.derivative: state_cx = horzcat(all_pn.nlp.states.cx_end, all_pn.nlp.states.cx) control_cx = horzcat(all_pn.nlp.controls.cx_end, all_pn.nlp.controls.cx) self.function = biorbd.to_casadi_func( f"{name}", self.function(all_pn.nlp.states.cx_end, all_pn.nlp.controls.cx_end, param_cx) - self.function(all_pn.nlp.states.cx, all_pn.nlp.controls.cx, param_cx), state_cx, control_cx, param_cx, ) modified_fcn = self.function(state_cx, control_cx, param_cx) dt_cx = nlp.cx.sym("dt", 1, 1) weight_cx = nlp.cx.sym("weight", 1, 1) target_cx = nlp.cx.sym("target", modified_fcn.shape) modified_fcn = modified_fcn - target_cx if self.weight: modified_fcn = modified_fcn ** 2 if self.quadratic else modified_fcn modified_fcn = weight_cx * modified_fcn * dt_cx else: modified_fcn = modified_fcn * dt_cx # Do not use nlp.add_casadi_func because all of them must be registered self.weighted_function = Function( name, [state_cx, control_cx, param_cx, weight_cx, target_cx, dt_cx], [modified_fcn] ) self.weighted_function_non_threaded = self.weighted_function if ocp.n_threads > 1 and self.multi_thread and len(self.node_idx) > 1: self.function = self.function.map(len(self.node_idx), "thread", ocp.n_threads) self.weighted_function = self.weighted_function.map(len(self.node_idx), "thread", ocp.n_threads) else: self.multi_thread = False # Override the multi_threading, since only one node is optimized if self.expand: self.function = self.function.expand() self.weighted_function = self.weighted_function.expand()
def _set_penalty_function(self, all_pn: Union[PenaltyNodeList, list, tuple], fcn: Union[MX, SX]): """ Finalize the preparation of the penalty (setting function and weighted_function) Parameters ---------- all_pn: PenaltyNodeList The nodes fcn: Union[MX, SX] The value of the penalty function """ # Sanity checks if self.transition and self.explicit_derivative: raise ValueError("transition and explicit_derivative cannot be true simultaneously") if self.transition and self.derivative: raise ValueError("transition and derivative cannot be true simultaneously") if self.derivative and self.explicit_derivative: raise ValueError("derivative and explicit_derivative cannot be true simultaneously") def get_u(nlp, u: Union[MX, SX], dt: Union[MX, SX]): """ Get the control at a given time Parameters ---------- nlp: NonlinearProgram The nonlinear program u: Union[MX, SX] The control matrix dt: Union[MX, SX] The time a which control should be computed Returns ------- The control at a given time """ if nlp.control_type == ControlType.CONSTANT: return u elif nlp.control_type == ControlType.LINEAR_CONTINUOUS: return u[:, 0] + (u[:, 1] - u[:, 0]) * dt else: raise RuntimeError(f"{nlp.control_type} ControlType not implemented yet") return u if self.multinode_constraint or self.transition: ocp = all_pn[0].ocp nlp = all_pn[0].nlp nlp_post = all_pn[1].nlp name = self.name.replace("->", "_").replace(" ", "_").replace(",", "_") states_pre = nlp.states.cx_end states_post = nlp_post.states.cx controls_pre = nlp.controls.cx_end controls_post = nlp_post.controls.cx state_cx = vertcat(states_pre, states_post) control_cx = vertcat(controls_pre, controls_post) else: ocp = all_pn.ocp nlp = all_pn.nlp name = self.name if self.integrate: state_cx = horzcat(*([all_pn.nlp.states.cx] + all_pn.nlp.states.cx_intermediates_list)) control_cx = all_pn.nlp.controls.cx else: state_cx = all_pn.nlp.states.cx control_cx = all_pn.nlp.controls.cx if self.explicit_derivative: if self.derivative: raise RuntimeError("derivative and explicit_derivative cannot be simultaneously true") state_cx = horzcat(state_cx, all_pn.nlp.states.cx_end) control_cx = horzcat(control_cx, all_pn.nlp.controls.cx_end) param_cx = nlp.cx(nlp.parameters.cx) # Do not use nlp.add_casadi_func because all functions must be registered sub_fcn = fcn[self.rows, self.cols] self.function = biorbd.to_casadi_func(name, sub_fcn, state_cx, control_cx, param_cx, expand=self.expand) self.function_non_threaded = self.function if self.derivative: state_cx = horzcat(all_pn.nlp.states.cx_end, all_pn.nlp.states.cx) control_cx = horzcat(all_pn.nlp.controls.cx_end, all_pn.nlp.controls.cx) self.function = biorbd.to_casadi_func( f"{name}", self.function(all_pn.nlp.states.cx_end, all_pn.nlp.controls.cx_end, param_cx) - self.function(all_pn.nlp.states.cx, all_pn.nlp.controls.cx, param_cx), state_cx, control_cx, param_cx, ) dt_cx = nlp.cx.sym("dt", 1, 1) is_trapezoidal = ( self.integration_rule == IntegralApproximation.TRAPEZOIDAL or self.integration_rule == IntegralApproximation.TRUE_TRAPEZOIDAL ) target_shape = tuple( [ len(self.rows), len(self.cols) + 1 if is_trapezoidal else len(self.cols), ] ) target_cx = nlp.cx.sym("target", target_shape) weight_cx = nlp.cx.sym("weight", 1, 1) exponent = 2 if self.quadratic and self.weight else 1 if is_trapezoidal: # Hypothesis: the function is continuous on states # it neglects the discontinuities at the beginning of the optimization state_cx = ( horzcat(all_pn.nlp.states.cx, all_pn.nlp.states.cx_end) if self.integration_rule == IntegralApproximation.TRAPEZOIDAL else all_pn.nlp.states.cx ) # to handle piecewise constant in controls we have to compute the value for the end of the interval # which only relies on the value of the control at the beginning of the interval control_cx = ( horzcat(all_pn.nlp.controls.cx) if nlp.control_type == ControlType.CONSTANT else horzcat(all_pn.nlp.controls.cx, all_pn.nlp.controls.cx_end) ) control_cx_end = get_u(nlp, control_cx, dt_cx) state_cx_end = ( all_pn.nlp.states.cx_end if self.integration_rule == IntegralApproximation.TRAPEZOIDAL else nlp.dynamics[0](x0=state_cx, p=control_cx_end, params=nlp.parameters.cx)["xf"] ) self.modified_function = biorbd.to_casadi_func( f"{name}", ( (self.function(all_pn.nlp.states.cx, all_pn.nlp.controls.cx, param_cx) - target_cx[:, 0]) ** exponent + (self.function(state_cx_end, control_cx_end, param_cx) - target_cx[:, 1]) ** exponent ) / 2, state_cx, control_cx, param_cx, target_cx, dt_cx, ) modified_fcn = self.modified_function(state_cx, control_cx, param_cx, target_cx, dt_cx) else: modified_fcn = (self.function(state_cx, control_cx, param_cx) - target_cx) ** exponent modified_fcn = weight_cx * modified_fcn * dt_cx if self.weight else modified_fcn * dt_cx # Do not use nlp.add_casadi_func because all of them must be registered self.weighted_function = Function( name, [state_cx, control_cx, param_cx, weight_cx, target_cx, dt_cx], [modified_fcn] ) self.weighted_function_non_threaded = self.weighted_function if ocp.n_threads > 1 and self.multi_thread and len(self.node_idx) > 1: self.function = self.function.map(len(self.node_idx), "thread", ocp.n_threads) self.weighted_function = self.weighted_function.map(len(self.node_idx), "thread", ocp.n_threads) else: self.multi_thread = False # Override the multi_threading, since only one node is optimized if self.expand: self.function = self.function.expand() self.weighted_function = self.weighted_function.expand()
def find_initial_root_pose(self): # This method finds a root pose such that the body of a given pose has its CoM centered to the feet model = self.models[0] n_root = model.nbRoot() body_pose_no_root = self.q_mapping.to_second.map( self.body_at_first_node)[n_root:, 0] bimap = BiMapping( list(range(n_root)) + [None] * body_pose_no_root.shape[0], list(range(n_root))) bound_min = [] bound_max = [] for i in range(model.nbSegment()): seg = model.segment(i) for r in seg.QRanges(): bound_min.append(r.min()) bound_max.append(r.max()) bound_min = bimap.to_first.map(np.array(bound_min)[:, np.newaxis]) bound_max = bimap.to_first.map(np.array(bound_max)[:, np.newaxis]) root_bounds = (list(bound_min[:, 0]), list(bound_max[:, 0])) q_sym = MX.sym("Q", model.nbQ(), 1) com_func = biorbd.to_casadi_func("com", model.CoM, q_sym) contacts_func = biorbd.to_casadi_func("contacts", model.constraintsInGlobal, q_sym, True) shoulder_jcs_func = biorbd.to_casadi_func("shoulder_jcs", model.globalJCS, q_sym, 3) hand_marker_func = biorbd.to_casadi_func("hand_marker", model.marker, q_sym, 32) def objective_function(q_root, *args, **kwargs): # Center of mass q = bimap.to_second.map(q_root[:, np.newaxis])[:, 0] q[model.nbRoot():] = body_pose_no_root com = np.array(com_func(q)) contacts = np.array(contacts_func(q)) mean_contacts = np.mean(contacts, axis=1) shoulder_jcs = np.array(shoulder_jcs_func(q)) hand = np.array(hand_marker_func(q)) # Prepare output out = np.ndarray((0, )) # The center of contact points should be at 0 out = np.concatenate((out, mean_contacts[0:2])) out = np.concatenate((out, contacts[2, self.flatfoot_contact_z_idx])) # The projection of the center of mass should be at 0 and at 0.95 meter high out = np.concatenate((out, (com + np.array([[0, 0, -0.95]]).T)[:, 0])) # Keep the arms horizontal out = np.concatenate((out, (shoulder_jcs[2, 3] - hand[2]))) return out q_root0 = np.mean(root_bounds, axis=0) pos = optimize.least_squares(objective_function, x0=q_root0, bounds=root_bounds) root = np.zeros(n_root) root[bimap.to_first.map_idx] = pos.x return root
def main(): """ Firstly, it solves the getting_started/pendulum.py example. Afterward, it gets the marker positions and joint torque from the solution and uses them to track. It then creates and solves this ocp and show the results """ biorbd_path = str( EXAMPLES_FOLDER) + "/getting_started/models/pendulum.bioMod" biorbd_model = biorbd.Model(biorbd_path) final_time = 1 n_shooting = 20 ocp_to_track = data_to_track.prepare_ocp(biorbd_model_path=biorbd_path, final_time=final_time, n_shooting=n_shooting) sol = ocp_to_track.solve() q, qdot, tau = sol.states["q"], sol.states["qdot"], sol.controls["tau"] n_q = biorbd_model.nbQ() n_marker = biorbd_model.nbMarkers() x = np.concatenate((q, qdot)) symbolic_states = MX.sym("q", n_q, 1) markers_fun = biorbd.to_casadi_func("ForwardKin", biorbd_model.markers, symbolic_states) markers_ref = np.zeros((3, n_marker, n_shooting + 1)) for i in range(n_shooting + 1): markers_ref[:, :, i] = markers_fun(x[:n_q, i]) tau_ref = tau[:, :-1] ocp = prepare_ocp( biorbd_model, final_time=final_time, n_shooting=n_shooting, markers_ref=markers_ref, tau_ref=tau_ref, ) # --- plot markers position --- # title_markers = ["x", "y", "z"] marker_color = ["tab:red", "tab:orange"] ocp.add_plot( "Markers plot coordinates", update_function=lambda t, x, u, p: get_markers_pos( x, 0, markers_fun, n_q), linestyle=".-", plot_type=PlotType.STEP, color=marker_color[0], ) ocp.add_plot( "Markers plot coordinates", update_function=lambda t, x, u, p: get_markers_pos( x, 1, markers_fun, n_q), linestyle=".-", plot_type=PlotType.STEP, color=marker_color[1], ) ocp.add_plot( "Markers plot coordinates", update_function=lambda t, x, u, p: markers_ref[:, 0, :], plot_type=PlotType.PLOT, color=marker_color[0], legend=title_markers, ) ocp.add_plot( "Markers plot coordinates", update_function=lambda t, x, u, p: markers_ref[:, 1, :], plot_type=PlotType.PLOT, color=marker_color[1], legend=title_markers, ) # --- Solve the program --- # sol = ocp.solve(Solver.IPOPT(show_online_optim=True)) # --- Show results --- # sol.animate(n_frames=100)