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)
