def shift_knot1_fwd(cfs, basis, t_shift):
if isinstance(cfs, (SX, MX)):
cfs_sym = MX.sym('cfs', cfs.shape)
t_shift_sym = MX.sym('t_shift')
T = shiftfirstknot_T(basis, t_shift_sym)
cfs2_sym = mtimes(T, cfs_sym)
fun = Function('fun', [cfs_sym, t_shift_sym], [cfs2_sym]).expand()
return fun(cfs, t_shift)
else:
T = shiftfirstknot_T(basis, t_shift)
return T.dot(cfs)
def construct_upd_z(self, problem=None, lineq_updz=True):
if problem is not None:
self.problem_upd_z = problem
self._lineq_updz = lineq_updz
return 0.
# check if we have linear equality constraints
self._lineq_updz, A, b = self._check_for_lineq()
if not self._lineq_updz:
raise ValueError('For now, only equality constrained QP ' +
'z-updates are allowed!')
x_i = struct_symMX(self.q_i_struct)
x_j = struct_symMX(self.q_ij_struct)
l_i = struct_symMX(self.q_i_struct)
l_ij = struct_symMX(self.q_ij_struct)
t = MX.sym('t')
T = MX.sym('T')
rho = MX.sym('rho')
par = struct_symMX(self.par_global_struct)
inp = [x_i.cat, l_i.cat, l_ij.cat, x_j.cat, t, T, rho, par.cat]
t0 = t/T
# put symbols in MX structs (necessary for transformation)
x_i = self.q_i_struct(x_i)
x_j = self.q_ij_struct(x_j)
l_i = self.q_i_struct(l_i)
l_ij = self.q_ij_struct(l_ij)
# transform spline variables: only consider future piece of spline
tf = lambda cfs, basis: shift_knot1_fwd(cfs, basis, t0)
self._transform_spline([x_i, l_i], tf, self.q_i)
self._transform_spline([x_j, l_ij], tf, self.q_ij)
# fill in parameters
A = A(par.cat)
b = b(par.cat)
# build KKT system and solve it via schur complement method
l, x = vertcat(l_i.cat, l_ij.cat), vertcat(x_i.cat, x_j.cat)
f = -(l + rho*x)
G = -(1/rho)*mtimes(A, A.T)
h = b + (1/rho)*mtimes(A, f)
mu = solve(G, h)
z = -(1/rho)*(mtimes(A.T, mu)+f)
l_qi = self.q_i_struct.shape[0]
l_qij = self.q_ij_struct.shape[0]
z_i_new = self.q_i_struct(z[:l_qi])
z_ij_new = self.q_ij_struct(z[l_qi:l_qi+l_qij])
# transform back
tf = lambda cfs, basis: shift_knot1_bwd(cfs, basis, t0)
self._transform_spline(z_i_new, tf, self.q_i)
self._transform_spline(z_ij_new, tf, self.q_ij)
out = [z_i_new.cat, z_ij_new.cat]
# create problem
prob, buildtime = create_function('upd_z_'+str(self._index), inp, out, self.options)
self.problem_upd_z = prob
return buildtime
def new_variable(self, name, dim, init, lb, ub):
"""
Add a new variable to the problem.
Parameters
----------
name : string
Variable name.
dim : int
Number of dimensions.
init : array, shape=(dim,)
Initial values.
lb : array, shape=(dim,)
Vector of lower bounds on variable values.
ub : array, shape=(dim,)
Vector of upper bounds on variable values.
"""
assert len(init) == len(lb) == len(ub) == dim
var = MX.sym(name, dim)
self.var_symbols.append(var)
self.initvals += init
self.var_index[name] = len(self.var_lbounds)
self.var_lbounds += lb
self.var_ubounds += ub
return var
def convert_array_to_external_forces(all_f_ext):
if not isinstance(all_f_ext, (list, tuple)):
raise RuntimeError(
"f_ext should be a list of (6 x nb_external_forces x nb_shooting) or (6 x nb_shooting) matrix"
)
sv_over_all_phases = []
for f_ext in all_f_ext:
f_ext = np.array(f_ext)
if len(f_ext.shape) < 2 or len(f_ext.shape) > 3:
raise RuntimeError(
"f_ext should be a list of (6 x nb_external_forces x nb_shooting) or (6 x nb_shooting) matrix"
)
if len(f_ext.shape) == 2:
f_ext = f_ext[:, :, np.newaxis]
if f_ext.shape[0] != 6:
raise RuntimeError(
"f_ext should be a list of (6 x nb_external_forces x nb_shooting) or (6 x nb_shooting) matrix"
)
sv_over_phase = []
for node in range(f_ext.shape[2]):
sv = biorbd.VecBiorbdSpatialVector()
for idx in range(f_ext.shape[1]):
sv.append(biorbd.SpatialVector(MX(f_ext[:, idx, node])))
sv_over_phase.append(sv)
sv_over_all_phases.append(sv_over_phase)
return sv_over_all_phases
def blockdiag(*matrices_list):
"""Receives a list of matrices and return a block diagonal.
:param DM|MX|SX matrices_list: list of matrices
"""
size_1 = sum([m.size1() for m in matrices_list])
size_2 = sum([m.size2() for m in matrices_list])
matrix_types = [type(m) for m in matrices_list]
if SX in matrix_types and MX in matrix_types:
raise ValueError(
"Can not mix MX and SX types. Types give: {}".format(matrix_types))
if SX in matrix_types:
matrix = SX.zeros(size_1, size_2)
elif MX in matrix_types:
matrix = MX.zeros(size_1, size_2)
else:
matrix = DM.zeros(size_1, size_2)
index_1 = 0
index_2 = 0
for m in matrices_list:
matrix[index_1:index_1 + m.size1(), index_2:index_2 + m.size2()] = m
index_1 += m.size1()
index_2 += m.size2()
return matrix
def custom_dynamic(states: MX, controls: MX, parameters: MX,
nlp: NonLinearProgram) -> tuple:
"""
The dynamics of the system using an external force (see custom_dynamics for more explanation)
Parameters
----------
states: MX
The current states of the system
controls: MX
The current controls of the system
parameters: MX
The current parameters of the system
nlp: NonLinearProgram
A reference to the phase of the ocp
Returns
-------
The state derivative
"""
q = DynamicsFunctions.get(nlp.states["q"], states)
qdot = DynamicsFunctions.get(nlp.states["qdot"], states)
tau = DynamicsFunctions.get(nlp.controls["tau"], controls)
force_vector = MX.zeros(6)
force_vector[5] = 100 * q[0]**2
f_ext = biorbd.VecBiorbdSpatialVector()
f_ext.append(biorbd.SpatialVector(force_vector))
qddot = nlp.model.ForwardDynamics(q, qdot, tau, f_ext).to_mx()
return qdot, qddot
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 forward_dynamics_torque_driven(states, controls, parameters, nlp):
"""
Forward dynamics (q, qdot, qddot -> tau) with external forces driven by joint torques (controls).
:param states: States. (MX.sym from CasADi)
:param controls: Controls. (MX.sym from CasADi)
:param nlp: An OptimalControlProgram class.
:param parameters: The MX associated to the parameters
:return: Vertcat of derived states. (MX.sym from CasADi)
"""
DynamicsFunctions.apply_parameters(parameters, nlp)
q, qdot, tau = DynamicsFunctions.dispatch_q_qdot_tau_data(
states, controls, nlp)
q_dot = nlp.model.computeQdot(q, qdot).to_mx()
qdot_reduced = nlp.mapping["q"].reduce.map(q_dot)
if nlp.external_forces:
dxdt = MX(nlp.nx, nlp.ns)
for i, f_ext in enumerate(nlp.external_forces):
qddot = nlp.model.ForwardDynamics(q, qdot, tau, f_ext).to_mx()
qddot_reduced = nlp.mapping["q_dot"].reduce.map(qddot)
dxdt[:, i] = vertcat(qdot_reduced, qddot_reduced)
else:
qddot = nlp.model.ForwardDynamics(q, qdot, tau).to_mx()
qddot_reduced = nlp.mapping["q_dot"].reduce.map(qddot)
dxdt = vertcat(qdot_reduced, qddot_reduced)
return dxdt
def _expr_apply(self, expr, **kwargs):
"""
Substitute placeholder symbols with actual decision variables,
or expressions involving decision variables
"""
subst_from, subst_to = self._get_subst_set(**kwargs)
return substitute([MX(expr)], subst_from, subst_to)[0]
def __init__(self):
self.elements: list = []
self.fake_elements: list = []
self._cx: Union[MX, SX, np.ndarray] = np.array([])
self._cx_end: Union[MX, SX, np.ndarray] = np.array([])
self._cx_intermediates: list = []
self.mx_reduced: MX = MX.sym("var", 0, 0)
def append(self, name: str, cx: list, mx: MX, bimapping: BiMapping):
"""
Add a new variable to the list
Parameters
----------
name: str
The name of the variable
cx: list
The list of SX or MX variable associated with this variable
mx: MX
The MX variable associated with this variable
bimapping: BiMapping
The Mapping of the MX against CX
"""
index = range(self._cx.shape[0], self._cx.shape[0] + cx[0].shape[0])
self._cx = vertcat(self._cx, cx[0])
self._cx_end = vertcat(self._cx_end, cx[-1])
for i, c in enumerate(cx[1:-1]):
if i >= len(self._cx_intermediates):
self._cx_intermediates.append(c)
else:
self._cx_intermediates[i] = vertcat(self._cx_intermediates[i],
c)
self.mx_reduced = vertcat(self.mx_reduced, MX.sym("var", cx[0].shape))
self.elements.append(
OptimizationVariable(name, mx, index, bimapping, self))
def forward_dynamics_torque_driven(states, controls, nlp):
"""
:param states: MX.sym from CasADi.
:param controls: MX.sym from CasADi.
:param nlp: An OptimalControlProgram class
:return: Vertcat of derived states.
"""
q, qdot, tau = Dynamics.__dispatch_q_qdot_tau_data(
states, controls, nlp)
qdot_reduced = nlp["q_mapping"].reduce.map(qdot)
if "external_forces" in nlp:
dxdt = MX(nlp["nx"], nlp["ns"])
for i, f_ext in enumerate(nlp["external_forces"]):
qddot = biorbd.Model.ForwardDynamics(nlp["model"], q, qdot,
tau, f_ext).to_mx()
qddot_reduced = nlp["q_dot_mapping"].reduce.map(qddot)
dxdt[:, i] = vertcat(qdot_reduced, qddot_reduced)
else:
qddot = biorbd.Model.ForwardDynamics(nlp["model"], q, qdot,
tau).to_mx()
qddot_reduced = nlp["q_dot_mapping"].reduce.map(qddot)
dxdt = vertcat(qdot_reduced, qddot_reduced)
return dxdt
def test_set_scalar(brbd):
def check_value(target):
if brbd.currentLinearAlgebraBackend() == 1:
assert m.segment(0).characteristics().mass().to_mx() == target
else:
assert m.segment(0).characteristics().mass() == target
m = brbd.Model("../../models/pyomecaman.bioMod")
m.segment(0).characteristics().setMass(10)
check_value(10)
m.segment(0).characteristics().setMass(11.0)
check_value(11.0)
with pytest.raises(ValueError, match="Scalar must be a 1x1 array or a float"):
m.segment(0).characteristics().setMass(np.array([]))
m.segment(0).characteristics().setMass(np.array((12,)))
check_value(12.0)
m.segment(0).characteristics().setMass(np.array([[13]]))
check_value(13.0)
with pytest.raises(ValueError, match="Scalar must be a 1x1 array or a float"):
m.segment(0).characteristics().setMass(np.array([[[14]]]))
if brbd.currentLinearAlgebraBackend() == 1:
from casadi import MX
m.segment(0).characteristics().setMass(MX(15))
check_value(15.0)
def intg_rk(self, f, X, U, P, Z):
assert Z.is_empty()
DT = MX.sym("DT")
t0 = MX.sym("t0")
# A single Runge-Kutta 4 step
k1 = f(x=X, u=U, p=P, t=t0, z=Z)
k2 = f(x=X + DT / 2 * k1["ode"], u=U, p=P, t=t0+DT/2, z=Z)
k3 = f(x=X + DT / 2 * k2["ode"], u=U, p=P, t=t0+DT/2)
k4 = f(x=X + DT * k3["ode"], u=U, p=P, t=t0+DT)
f0 = k1["ode"]
f1 = 2/DT*(k2["ode"]-k1["ode"])/2
f2 = 4/DT**2*(k3["ode"]-k2["ode"])/6
f3 = 4*(k4["ode"]-2*k3["ode"]+k1["ode"])/DT**3/24
poly_coeff = hcat([X, f0, f1, f2, f3])
return Function('F', [X, U, t0, DT, P], [X + DT / 6 * (k1["ode"] + 2 * k2["ode"] + 2 * k3["ode"] + k4["ode"]), poly_coeff, DT / 6 * (k1["quad"] + 2 * k2["quad"] + 2 * k3["quad"] + k4["quad"])], ['x0', 'u', 't0', 'DT', 'p'], ['xf', 'poly_coeff', 'qf'])
def append_faked_optim_var(name, optim_var, keys: list):
"""
Add a fake optim var by combining vars in keys
Parameters
----------
optim_var: OptimizationVariableList
states or controls
keys: list
The list of keys to combine
"""
index = []
mx = MX()
to_second = []
to_first = []
for key in keys:
index.extend(list(optim_var[key].index))
mx = vertcat(mx, optim_var[key].mx)
to_second.extend(
list(
np.array(optim_var[key].mapping.to_second.map_idx) +
len(to_second)))
to_first.extend(
list(
np.array(optim_var[key].mapping.to_first.map_idx) +
len(to_first)))
optim_var.append_fake(name, index, mx, BiMapping(to_second, to_first))
def test_include_constraint_equality_w_external_var(self):
theta = MX.sym('theta')
aop = AbstractOptimizationProblem()
x = aop.create_variable('x', 2)
aop.include_constraint(x + 2 == theta)
def contact_forces(nlp: NonLinearProgram, q, qdot, tau):
"""
Easy accessor for the contact forces in contact dynamics
Parameters
----------
nlp: NonLinearProgram
The phase of the program
q: Union[MX, SX]
The value of q from "get"
qdot: Union[MX, SX]
The value of qdot from "get"
tau: Union[MX, SX]
The value of tau from "get"
Returns
-------
The contact forces
"""
cs = nlp.model.getConstraints()
if nlp.external_forces:
all_cs = MX()
for i, f_ext in enumerate(nlp.external_forces):
nlp.model.ForwardDynamicsConstraintsDirect(
q, qdot, tau, cs, f_ext).to_mx()
all_cs = horzcat(all_cs, cs.getForce().to_mx())
return all_cs
else:
nlp.model.ForwardDynamicsConstraintsDirect(q, qdot, tau,
cs).to_mx()
return cs.getForce().to_mx()
def xia_model_configuration(ocp, nlp):
Problem.configure_q_qdot(nlp, True, False)
Problem.configure_tau(nlp, False, True)
Problem.configure_muscles(nlp, False, True)
x = MX()
for i in range(nlp["nbMuscle"]):
x = vertcat(x, MX.sym(f"Muscle_{nlp['muscleNames']}_active_{nlp['phase_idx']}", 1, 1))
for i in range(nlp["nbMuscle"]):
x = vertcat(x, MX.sym(f"Muscle_{nlp['muscleNames']}_fatigue_{nlp['phase_idx']}", 1, 1))
for i in range(nlp["nbMuscle"]):
x = vertcat(x, MX.sym(f"Muscle_{nlp['muscleNames']}_resting_{nlp['phase_idx']}", 1, 1))
nlp["x"] = vertcat(nlp["x"], x)
nlp["var_states"]["muscles_active"] = nlp["nbMuscle"]
nlp["var_states"]["muscles_fatigue"] = nlp["nbMuscle"]
nlp["var_states"]["muscles_resting"] = nlp["nbMuscle"]
nlp["nx"] = nlp["x"].rows()
nb_q_qdot = nlp["nbQ"] + nlp["nbQdot"]
nlp["plot"]["muscles_active"] = CustomPlot(
lambda x, u, p: x[nb_q_qdot : nb_q_qdot + nlp["nbMuscle"]],
plot_type=PlotType.INTEGRATED,
legend=nlp["muscleNames"],
color="r",
ylim=[0, 1],
)
combine = "muscles_active"
nlp["plot"]["muscles_fatigue"] = CustomPlot(
lambda x, u, p: x[nb_q_qdot + nlp["nbMuscle"] : nb_q_qdot + 2 * nlp["nbMuscle"]],
plot_type=PlotType.INTEGRATED,
legend=nlp["muscleNames"],
combine_to=combine,
color="g",
ylim=[0, 1],
)
nlp["plot"]["muscles_resting"] = CustomPlot(
lambda x, u, p: x[nb_q_qdot + 2 * nlp["nbMuscle"] : nb_q_qdot + 3 * nlp["nbMuscle"]],
plot_type=PlotType.INTEGRATED,
legend=nlp["muscleNames"],
combine_to=combine,
color="b",
ylim=[0, 1],
)
Problem.configure_forward_dyn_func(ocp, nlp, xia_model_dynamic)
def configure_tau(nlp, as_states, as_controls):
"""
Configures common settings for torque driven problems with and without contacts.
:param nlp: An OptimalControlProgram class.
"""
if nlp["tau_mapping"] is None:
nlp["tau_mapping"] = BidirectionalMapping(
# Mapping(range(nlp["model"].nbGeneralizedTorque())), Mapping(range(nlp["model"].nbGeneralizedTorque()))
Mapping(range(nlp["model"].nbQdot())),
Mapping(
range(nlp["model"].nbQdot())
), # To change when nlp["model"].nbGeneralizedTorque() will return the proper number
)
dof_names = nlp["model"].nameDof()
n_col = nlp["control_type"].value
tau_mx = MX()
all_tau = [nlp["CX"]() for _ in range(n_col)]
for i in nlp["tau_mapping"].reduce.map_idx:
for j in range(len(all_tau)):
all_tau[j] = vertcat(all_tau[j], nlp["CX"].sym(f"Tau_{dof_names[i].to_string()}_{j}", 1, 1))
for i in nlp["q_mapping"].expand.map_idx:
tau_mx = vertcat(tau_mx, MX.sym("Tau_" + dof_names[i].to_string(), 1, 1))
nlp["nbTau"] = nlp["tau_mapping"].reduce.len
legend_tau = ["tau_" + nlp["model"].nameDof()[idx].to_string() for idx in nlp["tau_mapping"].reduce.map_idx]
nlp["tau"] = tau_mx
if as_states:
nlp["x"] = vertcat(nlp["x"], all_tau[0])
nlp["var_states"]["tau"] = nlp["nbTau"]
# Add plot if it happens
if as_controls:
nlp["u"] = vertcat(nlp["u"], horzcat(*all_tau))
nlp["var_controls"]["tau"] = nlp["nbTau"]
if nlp["control_type"] == ControlType.LINEAR_CONTINUOUS:
plot_type = PlotType.PLOT
else:
plot_type = PlotType.STEP
nlp["plot"]["tau"] = CustomPlot(lambda x, u, p: u[: nlp["nbTau"]], plot_type=plot_type, legend=legend_tau)
nlp["nx"] = nlp["x"].rows()
nlp["nu"] = nlp["u"].rows()
def test_is_equality(self):
x_vector = vertcat(MX.sym('x'), MX.sym('other_name'),
MX.sym('the_right_one'), MX.sym('foo'),
MX.sym('food'), MX.sym('bar'), MX.sym('var1'),
MX.sym('alpha'), MX.sym('cat'))
self.assertTrue(is_equality(x_vector[1] == 1))
self.assertFalse(is_equality(x_vector[1] >= 1))
self.assertFalse(is_equality(x_vector[1] <= 1))
self.assertTrue(is_equality(x_vector == 1))
self.assertFalse(is_equality(x_vector >= 1))
self.assertFalse(is_equality(x_vector <= 1))
def test_include_variable_with_bounds(self):
lb = -DM(range(1, 4))
ub = DM(range(5, 8))
aop = AbstractOptimizationProblem()
x = MX.sym('x', 3)
aop.include_variable(x, lb=lb, ub=ub)
self.assertTrue(is_equal(aop.x_lb, lb))
self.assertTrue(is_equal(aop.x_ub, ub))
def test_forward_dynamics():
m = biorbd.Model("../../models/pyomecaman_withActuators.bioMod")
q = np.array([i * 1.1 for i in range(m.nbQ())])
qdot = np.array([i * 1.1 for i in range(m.nbQ())])
tau = np.array([i * 1.1 for i in range(m.nbQ())])
if biorbd.currentLinearAlgebraBackend() == 1:
# If CasADi backend is used
from casadi import Function, MX
q_sym = MX.sym("q", m.nbQ(), 1)
qdot_sym = MX.sym("qdot", m.nbQdot(), 1)
tau_sym = MX.sym("tau", m.nbGeneralizedTorque(), 1)
ForwardDynamics = Function(
"ForwardDynamics",
[q_sym, qdot_sym, tau_sym],
[m.ForwardDynamics(q_sym, qdot_sym, tau_sym).to_mx()],
["q_sym", "qdot_sym", "tau_sym"],
["qddot"],
).expand()
qddot = ForwardDynamics(q, qdot, tau)
qddot = np.array(qddot)[:, 0]
elif biorbd.currentLinearAlgebraBackend() == 0:
# if Eigen backend is used
qddot = m.ForwardDynamics(q, qdot, tau).to_array()
qddot_expected = np.array([
20.554883896960259,
-22.317642013324736,
-77.406439058256126,
17.382961188212313,
-63.426361095191858,
93.816468824985876,
106.46105024484631,
95.116641811710167,
-268.1961283528546,
2680.3632159799949,
-183.4582596257801,
755.89411812405604,
163.60239754283589,
])
np.testing.assert_almost_equal(qddot, qddot_expected)
def muscle_activations_driven(nlp):
"""
Names states (nlp.x) and controls (nlp.u) and gives size to (nlp.nx) and (nlp.nu).
Works with torques and muscles.
:param nlp: An OptimalControlProgram class.
"""
ProblemType.__configure_q_qdot(nlp, True, False)
ProblemType.__configure_muscles(nlp, False, True)
u = MX()
for i in range(nlp["nbMuscle"]):
u = vertcat(u, MX.sym(f"Muscle_{nlp['muscleNames']}_activation"))
nlp["u"] = vertcat(nlp["u"], u)
nlp["var_controls"] = {"muscles": nlp["nbMuscle"]}
ProblemType.__configure_forward_dyn_func(
nlp, Dynamics.forward_dynamics_muscle_activations_driven)
def markers_fun(biorbd_model):
qMX = MX.sym("qMX", biorbd_model.nbQ())
return Function("markers", [qMX], [
horzcat(*[
biorbd_model.markers(qMX)[i].to_mx()
for i in range(biorbd_model.nbMarkers())
])
])
def __init__(self, m):
self.model = m
self.q = MX.sym("q", m.nbQ(), 1)
self.casadi_func = {}
self.mapping = {
"q": BiMapping(range(self.model.nbQ()),
range(self.model.nbQ()))
}
def _define_mx(self, name, size0, size1, dictionary, value=None):
if value is None:
value = np.zeros((size0, size1))
symbol_name = self._add_label(name)
dictionary[name] = MX.sym(symbol_name, size0, size1)
self._values[name] = value
self.add_to_dict(dictionary[name], name)
return dictionary[name]
def _define_mx(self, name, size0, size1, dictionary, value=None):
if value is None:
value = np.zeros((size0, size1))
symbol_name = self._add_label(name)
dictionary[name] = MX.sym(symbol_name, size0, size1)
self._values[name] = value
self.add_to_dict(dictionary[name], name)
return dictionary[name]
def contact_force_continuity(pen_node: PenaltyNode, idx_pre, idx_post):
final_contact_z = sum1(pen_node[0].nlp.contact_forces_func(pen_node[0].x[0], pen_node[0].u[0], pen_node[0].p)[idx_pre, :])
if idx_post:
starting_contact_z = sum1(pen_node[1].nlp.contact_forces_func(pen_node[1].x[0], pen_node[1].u[0], pen_node[1].p)[idx_post, :])
else:
starting_contact_z = MX.zeros(final_contact_z.shape)
return final_contact_z - starting_contact_z
def muscle_activations_and_torque_driven(nlp):
"""
Names states (nlp.x) and controls (nlp.u) and gives size to (nlp.nx) and (nlp.nu).
Works with torques and muscles.
:param nlp: An OptimalControlProgram class.
"""
nlp["dynamics_func"] = Dynamics.forward_dynamics_torque_muscle_driven
ProblemType.__configure_torque_driven(nlp)
nlp["nbMuscle"] = nlp["model"].nbMuscleTotal()
u = MX()
muscle_names = nlp["model"].muscleNames()
for i in range(nlp["nbMuscle"]):
u = vertcat(u, MX.sym("Muscle_" + muscle_names[i].to_string() + "_activation"))
nlp["u"] = vertcat(nlp["u"], u)
nlp["nu"] = nlp["u"].rows()
def test_trigProp(before_trigProp):
m, v, idx, a, M, V, C = before_trigProp
m_cas = MX.sym('m', (3, 1))
v_cas = MX.sym('v_cas', (3, 3))
M_sym, V_sym, C_sym = utils_cas.trig_prop(m_cas, v_cas, idx, a)
f = Function('trigProp', [m_cas, v_cas], [M_sym, V_sym, C_sym])
M_out, V_out, C_out = f(m, v)
a_tol = 1e-3
assert np.allclose(M_out, M, atol=a_tol), 'Are the output means the same '
assert np.allclose(V_out, V,
atol=a_tol), 'Are the output variances the same'
assert np.allclose(C_out, C,
atol=a_tol), 'Are the cross-covariances the same'
def configure_q(nlp, as_states, as_controls):
"""
Configures common settings for torque driven problems with and without contacts.
:param nlp: An OptimalControlProgram class.
"""
if nlp.mapping["q"] is None:
nlp.mapping["q"] = BidirectionalMapping(
Mapping(range(nlp.model.nbQ())),
Mapping(range(nlp.model.nbQ())))
dof_names = nlp.model.nameDof()
q_mx = MX()
q = nlp.CX()
for i in nlp.mapping["q"].reduce.map_idx:
q = vertcat(q, nlp.CX.sym("Q_" + dof_names[i].to_string(), 1, 1))
for i in nlp.mapping["q"].expand.map_idx:
q_mx = vertcat(q_mx, MX.sym("Q_" + dof_names[i].to_string(), 1, 1))
nlp.shape["q"] = nlp.mapping["q"].reduce.len
legend_q = [
"q_" + nlp.model.nameDof()[idx].to_string()
for idx in nlp.mapping["q"].reduce.map_idx
]
nlp.q = q_mx
if as_states:
nlp.x = vertcat(nlp.x, q)
nlp.var_states["q"] = nlp.shape["q"]
q_bounds = nlp.x_bounds[:nlp.shape["q"]]
nlp.plot["q"] = CustomPlot(
lambda x, u, p: x[:nlp.shape["q"]],
plot_type=PlotType.INTEGRATED,
legend=legend_q,
bounds=q_bounds,
)
if as_controls:
nlp.u = vertcat(nlp.u, q)
nlp.var_controls["q"] = nlp.shape["q"]
# Add plot (and retrieving bounds if plots of bounds) if this problem is ever added
nlp.nx = nlp.x.rows()
nlp.nu = nlp.u.rows()
def parameter(self, n_rows=1, n_cols=1, grid=''):
"""
Create a parameter
"""
# Create a placeholder symbol with a dummy name (see #25)
p = MX.sym("p", n_rows, n_cols)
self.parameters[grid].append(p)
self._set_transcribed(False)
return p
def new_variable(self, name, dim, init, lb, ub):
assert len(init) == len(lb) == len(ub) == dim
var = MX.sym(name, dim)
self.var_symbols.append(var)
self.initvals += init
self.var_index[name] = len(self.var_lbounds)
self.var_lbounds += lb
self.var_ubounds += ub
return var
def construct_upd_l(self, problem=None):
if problem is not None:
self.problem_upd_l = problem
return 0.
# create parameters
z_ij = struct_symMX(self.q_ij_struct)
l_ij = struct_symMX(self.q_ij_struct)
x_j = struct_symMX(self.q_ij_struct)
t = MX.sym('t')
T = MX.sym('T')
rho = MX.sym('rho')
inp = [x_j, z_ij, l_ij, t, T, rho]
# update lambda
l_ij_new = self.q_ij_struct(l_ij.cat + rho*(x_j.cat - z_ij.cat))
out = [l_ij_new]
# create problem
prob, buildtime = create_function('upd_l_'+str(self._index), inp, out, self.options)
self.problem_upd_l = prob
return buildtime
def construct_upd_res(self, problem=None):
if problem is not None:
self.problem_upd_res = problem
return 0.
# create parameters
x_i = struct_symMX(self.q_i_struct)
z_i = struct_symMX(self.q_i_struct)
z_i_p = struct_symMX(self.q_i_struct)
z_ij = struct_symMX(self.q_ij_struct)
z_ij_p = struct_symMX(self.q_ij_struct)
x_j = struct_symMX(self.q_ij_struct)
t = MX.sym('t')
T = MX.sym('T')
t0 = t/T
rho = MX.sym('rho')
inp = [x_i, z_i, z_i_p, z_ij, z_ij_p, x_j, t, T, rho]
# put symbols in MX structs (necessary for transformation)
x_i = self.q_i_struct(x_i)
z_i = self.q_i_struct(z_i)
z_i_p = self.q_i_struct(z_i_p)
z_ij = self.q_ij_struct(z_ij)
z_ij_p = self.q_ij_struct(z_ij_p)
x_j = self.q_ij_struct(x_j)
# transform spline variables: only consider future piece of spline
tf = lambda cfs, basis: shift_knot1_fwd(cfs, basis, t0)
self._transform_spline([x_i, z_i, z_i_p], tf, self.q_i)
self._transform_spline([x_j, z_ij, z_ij_p], tf, self.q_ij)
# compute residuals
pr = mtimes((x_i.cat-z_i.cat).T, (x_i.cat-z_i.cat))
pr += mtimes((x_j.cat-z_ij.cat).T, (x_j.cat-z_ij.cat))
dr = rho*mtimes((z_i.cat-z_i_p.cat).T, (z_i.cat-z_i_p.cat))
dr += rho*mtimes((z_ij.cat-z_ij_p.cat).T, (z_ij.cat-z_ij_p.cat))
cr = rho*pr + dr
out = [pr, dr, cr]
# create problem
prob, buildtime = create_function('upd_res_'+str(self._index), inp, out, self.options)
self.problem_upd_res = prob
return buildtime
def _construct_upd_z_nlp(self):
# construct variables
self._var_struct_updz = struct([entry('z_i', struct=self.q_i_struct),
entry('z_ij', struct=self.q_ij_struct)])
var = struct_symMX(self._var_struct_updz)
z_i = self.q_i_struct(var['z_i'])
z_ij = self.q_ij_struct(var['z_ij'])
# construct parameters
self._par_struct_updz = struct([entry('x_i', struct=self.q_i_struct),
entry('x_j', struct=self.q_ij_struct),
entry('l_i', struct=self.q_i_struct),
entry('l_ij', struct=self.q_ij_struct),
entry('t'), entry('T'), entry('rho'),
entry('par', struct=self.par_struct)])
par = struct_symMX(self._par_struct_updz)
x_i, x_j = self.q_i_struct(par['x_i']), self.q_ij_struct(par['x_j'])
l_i, l_ij = self.q_i_struct(par['l_i']), self.q_ij_struct(par['l_ij'])
t, T, rho = par['t'], par['T'], par['rho']
t0 = t/T
# transform spline variables: only consider future piece of spline
tf = lambda cfs, basis: shift_knot1_fwd(cfs, basis, t0)
self._transform_spline([x_i, z_i, l_i], tf, self.q_i)
self._transform_spline([x_j, z_ij, l_ij], tf, self.q_ij)
# construct constraints
constraints, lb, ub = [], [], []
for con in self.constraints:
c = con[0]
for sym in symvar(c):
for label, child in self.group.items():
if sym.getName() in child.symbol_dict:
name = child.symbol_dict[sym.getName()][1]
v = z_i[label, name]
ind = self.q_i[child][name]
sym2 = MX.zeros(sym.size())
sym2[ind] = v
sym2 = reshape(sym2, sym.shape)
c = substitute(c, sym, sym2)
break
for nghb in self.q_ij.keys():
for label, child in nghb.group.items():
if sym.getName() in child.symbol_dict:
name = child.symbol_dict[sym.getName()][1]
v = z_ij[nghb.label, label, name]
ind = self.q_ij[nghb][child][name]
sym2 = MX.zeros(sym.size())
sym2[ind] = v
sym2 = reshape(sym2, sym.shape)
c = substitute(c, sym, sym2)
break
for name, s in self.par_i.items():
if s.getName() == sym.getName():
c = substitute(c, sym, par['par', name])
constraints.append(c)
lb.append(con[1])
ub.append(con[2])
self.lb_updz, self.ub_updz = lb, ub
# construct objective
obj = 0.
for child, q_i in self.q_i.items():
for name in q_i.keys():
x = x_i[child.label, name]
z = z_i[child.label, name]
l = l_i[child.label, name]
obj += mul(l.T, x-z) + 0.5*rho*mul((x-z).T, (x-z))
for nghb in self.q_ij.keys():
for child, q_ij in self.q_ij[nghb].items():
for name in q_ij.keys():
x = x_j[str(nghb), child.label, name]
z = z_ij[str(nghb), child.label, name]
l = l_ij[str(nghb), child.label, name]
obj += mul(l.T, x-z) + 0.5*rho*mul((x-z).T, (x-z))
# construct problem
prob, compile_time = self.father.create_nlp(var, par, obj,
constraints, self.options)
self.problem_upd_z = prob
def construct_upd_xz(self, problem=None):
# construct optifather & give reference to problem
self.father_updx = OptiFather(self.group.values())
self.problem.father = self.father_updx
# define z_ij variables
init = self.q_ij_struct(0)
for nghb, q_ij in self.q_ij.items():
for child, q_j in q_ij.items():
for name, ind in q_j.items():
var = np.array(child._values[name])
v = var.T.flatten()[ind]
init[nghb.label, child.label, name, ind] = v
z_ij = self.define_variable(
'z_ij', self.q_ij_struct.shape[0], value=np.array(init.cat))
# define parameters
l_ij = self.define_parameter('l_ij', self.q_ij_struct.shape[0])
l_ji = self.define_parameter('l_ji', self.q_ji_struct.shape[0])
# put them in the struct format
z_ij = self.q_ij_struct(z_ij)
l_ij = self.q_ij_struct(l_ij)
l_ji = self.q_ji_struct(l_ji)
# get (part of) variables
x_i = self._get_x_variables(symbolic=True)
# construct local copies of parameters
par = {}
for name, s in self.par_global.items():
par[name] = self.define_parameter(name, s.shape[0], s.shape[1])
if problem is None:
# get time info
t = self.define_symbol('t')
T = self.define_symbol('T')
t0 = t/T
# transform spline variables: only consider future piece of spline
tf = lambda cfs, basis: shift_knot1_fwd(cfs, basis, t0)
self._transform_spline(x_i, tf, self.q_i)
self._transform_spline([z_ij, l_ij], tf, self.q_ij)
self._transform_spline(l_ji, tf, self.q_ji)
# construct objective
obj = 0.
for child, q_i in self.q_i.items():
for name in q_i.keys():
x = x_i[child.label][name]
for nghb in self.q_ji.keys():
l = l_ji[str(nghb), child.label, name]
obj += mtimes(l.T, x)
for nghb, q_j in self.q_ij.items():
for child in q_j.keys():
for name in q_j[child].keys():
z = z_ij[str(nghb), child.label, name]
l = l_ij[str(nghb), child.label, name]
obj -= mtimes(l.T, z)
self.define_objective(obj)
# construct constraints
for con in self.global_constraints:
c = con[0]
for sym in symvar(c):
for label, child in self.group.items():
if sym.name() in child.symbol_dict:
name = child.symbol_dict[sym.name()][1]
v = x_i[label][name]
ind = self.q_i[child][name]
sym2 = MX.zeros(sym.size())
sym2[ind] = v
sym2 = reshape(sym2, sym.shape)
c = substitute(c, sym, sym2)
break
for nghb in self.q_ij.keys():
for label, child in nghb.group.items():
if sym.name() in child.symbol_dict:
name = child.symbol_dict[sym.name()][1]
v = z_ij[nghb.label, label, name]
ind = self.q_ij[nghb][child][name]
sym2 = MX.zeros(sym.size())
sym2[ind] = v
sym2 = reshape(sym2, sym.shape)
c = substitute(c, sym, sym2)
break
for name, s in self.par_global.items():
if s.name() == sym.name():
c = substitute(c, sym, par[name])
lb, ub = con[1], con[2]
self.define_constraint(c, lb, ub)
# construct problem
prob, buildtime = self.father_updx.construct_problem(
self.options, str(self._index), problem)
self.problem_upd_xz = prob
self.father_updx.init_transformations(self.problem.init_primal_transform,
self.problem.init_dual_transform)
self.init_var_dd()
return buildtime
