-
Notifications
You must be signed in to change notification settings - Fork 3
/
Copy pathmiller_viz.py
133 lines (111 loc) · 3.94 KB
/
miller_viz.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
"""
This file contains functions to plot extra information for the optimal control problem of the Miller.
"""
import numpy as np
import biorbd_casadi as biorbd
from bioptim import PlotType, NonLinearProgram, OptimalControlProgram
from custom_dynamics.enums import MillerDynamics
def plot_linear_momentum(x, nlp: NonLinearProgram):
"""
Compute the linear momentum of the system.
Parameters
----------
x
State vector
nlp: NonLinearProgram
Non linear program
"""
linear_momentum = (
nlp.model.mass().to_mx() * nlp.model.CoMdot(nlp.states["q"].mx, nlp.states["qdot"].mx, True).to_mx()
)
linear_momentum_func = biorbd.to_casadi_func(
"LinearMomentum", linear_momentum, 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.array(linear_momentum_func(q, qdot))
def plot_angular_momentum(x, nlp: NonLinearProgram):
"""
Compute the angular momentum of the system.
Parameters
----------
x
State vector
nlp: NonLinearProgram
Non linear program
"""
angular_momentum_func = biorbd.to_casadi_func(
"AngularMomentum", nlp.model.angularMomentum, 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.array(angular_momentum_func(q, qdot))
def plot_residual_torque(x, u, nlp: NonLinearProgram):
"""
Compute the residual torque of the system.
Parameters
----------
x
State vector
u
Control vector
nlp: NonLinearProgram
Non linear program
"""
qddot_mx = nlp.controls["qddot"].mx if "qddot" in nlp.controls.keys() else nlp.states["qddot"].mx
ID_func = biorbd.to_casadi_func(
"ResidualTorque",
nlp.model.InverseDynamics,
nlp.states["q"].mx,
nlp.states["qdot"].mx,
qddot_mx,
expand=True,
)
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, :])
qddot = (
u[nlp.controls["qddot"].index, :]
if "qddot" in nlp.controls.keys()
else nlp.states["qddot"].mapping.to_second.map(x[nlp.states["qddot"].index, :])
)
return np.array(ID_func(q, qdot, qddot)[:6, :])
def add_custom_plots(ocp: OptimalControlProgram, dynamics_type: MillerDynamics):
"""
Add extra plots to the OCP.
Parameters
----------
ocp: OptimalControlProgram
Optimal control program
dynamics_type: MillerDynamics
Type of dynamics of the Miller optimal control problem
"""
for i, nlp in enumerate(ocp.nlp):
ocp.add_plot(
"LinearMomentum",
lambda t, x, u, p: plot_linear_momentum(x, nlp),
phase=i,
legend=["Linear Momentum x", "Linear Momentum y", "Linear Momentum z"],
plot_type=PlotType.PLOT,
)
ocp.add_plot(
"AngularMomentum",
lambda t, x, u, p: plot_angular_momentum(x, nlp),
phase=i,
legend=["Angular Momentum x", "Angular Momentum y", "Angular Momentum z"],
plot_type=PlotType.PLOT,
)
if "implicit" in dynamics_type.value:
ocp.add_plot(
"TorqueResiduals",
lambda t, x, u, p: plot_residual_torque(x, u, nlp),
phase=i,
legend=[
"TorqueResiduals x",
"TorqueResiduals y",
"TorqueResiduals z",
"TorqueResiduals Rx",
"TorqueResiduals Ry",
"TorqueResiduals Rz",
],
plot_type=PlotType.PLOT,
)