api/bioptim_gui_api/acrobatics_ocp/variables/variable_computers/straight_acrobatics_variables.py
import numpy as np
from bioptim_gui_api.acrobatics_ocp.variables.utils import maximum_fig_arms_angle
from bioptim_gui_api.utils.format_utils import invert_min_max
from bioptim_gui_api.variables.misc.variables_utils import define_loose_bounds, LooseValue
class StraightAcrobaticsVariables:
X = 0
Y = 1
Z = 2
Xrot = 3
Yrot = 4
Zrot = 5
ZrotRightUpperArm = 6
YrotRightUpperArm = 7
ZrotLeftUpperArm = 8
YrotLeftUpperArm = 9
dofs = [
"Pelvis translation X",
"Pelvis translation Y",
"Pelvis translation Z",
"Pelvis rotation X",
"Pelvis rotation Y",
"Pelvis rotation Z",
"Right upper arm rotation Z",
"Right upper arm rotation Y",
"Left upper arm rotation Z",
"Left upper arm rotation Y",
]
nb_q, nb_qdot, nb_tau = 10, 10, 4
tau_min, tau_max, tau_init = -500.0, 500.0, 0.0
qdot_min, qdot_max = -10 * np.pi, 10 * np.pi
arm_dofs = [
ZrotRightUpperArm,
YrotRightUpperArm,
ZrotLeftUpperArm,
YrotLeftUpperArm,
]
shoulder_dofs = [
ZrotRightUpperArm,
YrotRightUpperArm,
ZrotLeftUpperArm,
YrotLeftUpperArm,
]
elbow_dofs = []
legs_xdofs = []
q_min_bounds = np.array(
[
[-1] * 3,
[-1] * 3,
[-0.1] * 3,
[0] * 3,
[-np.pi / 4] * 3,
[0] * 3,
[-0.65] * 3,
[-0.05] * 3,
[-2.0] * 3,
[-3.0] * 3,
]
)
q_max_bounds = np.array(
[
[1] * 3,
[1] * 3,
[15] * 3,
[0] * 3,
[np.pi / 4] * 3,
[0] * 3,
[2.0] * 3,
[3.0] * 3,
[0.65] * 3,
[0.05] * 3,
]
)
@classmethod
def _fill_init_phase(cls, x_bounds: np.ndarray) -> None:
x_bounds[0]["min"][:, 0] = [0] * cls.nb_q
x_bounds[0]["min"][: cls.Xrot, 0] = -0.001
x_bounds[0]["min"][[cls.YrotRightUpperArm, cls.YrotLeftUpperArm], 0] = 2.9, -2.9
x_bounds[0]["max"][:, 0] = [0] * cls.nb_q
x_bounds[0]["max"][: cls.Xrot, 0] = 0.001
x_bounds[0]["max"][[cls.YrotRightUpperArm, cls.YrotLeftUpperArm], 0] = 2.9, -2.9
@classmethod
def _fill_somersault_phase(cls, x_bounds: np.ndarray, phase: int, half_twists: list) -> None:
nb_somersaults = len(half_twists)
if phase != 0:
# initial bounds, same as final bounds of previous phase
x_bounds[phase]["min"][:, 0] = x_bounds[phase - 1]["min"][:, 2]
x_bounds[phase]["max"][:, 0] = x_bounds[phase - 1]["max"][:, 2]
intermediate_min_bounds = cls.q_min_bounds.copy()[:, 0]
intermediate_max_bounds = cls.q_max_bounds.copy()[:, 0]
# Intermediate bounds, same for every phase
x_bounds[phase]["min"][:, 1] = intermediate_min_bounds
x_bounds[phase]["max"][:, 1] = intermediate_max_bounds
x_bounds[phase]["min"][:, 2] = intermediate_min_bounds
x_bounds[phase]["max"][:, 2] = intermediate_max_bounds
# somersaulting
x_bounds[phase]["min"][cls.Xrot, 1] = 2 * np.pi * phase - 0.1
x_bounds[phase]["max"][cls.Xrot, 1] = 2 * np.pi * (phase + 1) + 0.1
define_loose_bounds(x_bounds[phase], cls.Xrot, 2, LooseValue(2 * np.pi * (phase + 1), 0.1))
# twisting
x_bounds[phase]["min"][cls.Zrot, 1] = np.pi * sum(half_twists[:phase]) - np.pi / 4 - 0.2
x_bounds[phase]["max"][cls.Zrot, 1] = np.pi * sum(half_twists[: phase + 1]) + np.pi / 4 + 0.2
define_loose_bounds(
x_bounds[phase], cls.Zrot, 2, LooseValue(np.pi * sum(half_twists[: phase + 1]), np.pi / 4 + 0.2)
)
if phase == nb_somersaults - 1:
# bounds for last_somersault
# keep 1/2 somersault before landing phase
x_bounds[nb_somersaults - 1]["min"][cls.Xrot, 1] = 2 * np.pi * (nb_somersaults - 1) - 0.1
x_bounds[nb_somersaults - 1]["max"][cls.Xrot, 1] = 2 * np.pi * nb_somersaults - np.pi / 2 + 0.1
define_loose_bounds(
x_bounds[nb_somersaults - 1], cls.Xrot, 2, LooseValue(2 * np.pi * nb_somersaults - np.pi / 2, 0.1)
)
# twists must be done before landing
define_loose_bounds(x_bounds[nb_somersaults - 1], cls.Zrot, 2, LooseValue(np.pi * sum(half_twists), 0.1))
@classmethod
def _fill_landing_phase(cls, x_bounds, half_twists: list) -> dict:
nb_somersaults = len(half_twists)
# Pelvis translations
x_bounds[-1]["min"][[cls.X, cls.Y, cls.Z], 2] = [-0.01, -0.01, 0]
x_bounds[-1]["max"][[cls.X, cls.Y, cls.Z], 2] = [0.01, 0.01, 0.01]
# Somersault
x_bounds[-1]["min"][cls.Xrot, 0] = x_bounds[-2]["min"][cls.Xrot, 2]
x_bounds[-1]["max"][cls.Xrot, 0] = x_bounds[-2]["max"][cls.Xrot, 2]
x_bounds[-1]["min"][cls.Xrot, 1] = 2 * np.pi * nb_somersaults - np.pi / 2 - 0.1
x_bounds[-1]["max"][cls.Xrot, 1] = 2 * np.pi * nb_somersaults + 0.1
define_loose_bounds(x_bounds[-1], cls.Xrot, 2, LooseValue(2 * np.pi * nb_somersaults, 0.1))
# Tilt
define_loose_bounds(x_bounds[-1], cls.Yrot, None, LooseValue(0.0, np.pi / 16))
# Twist
define_loose_bounds(x_bounds[-1], cls.Zrot, None, LooseValue(np.pi * sum(half_twists), 0.1))
# FIG Code of Points 14.5, arms to stop twisting rotation
max_angle = maximum_fig_arms_angle(half_twists)
# Right arm
x_bounds[-1]["min"][cls.YrotRightUpperArm, 0] = 0
x_bounds[-1]["max"][cls.YrotRightUpperArm, 0] = max_angle
x_bounds[-2]["min"][cls.YrotRightUpperArm, 2] = 0
x_bounds[-2]["max"][cls.YrotRightUpperArm, 2] = max_angle
define_loose_bounds(x_bounds[-1], cls.YrotRightUpperArm, 2, LooseValue(2.9, 0.1))
define_loose_bounds(x_bounds[-1], cls.ZrotRightUpperArm, 2, LooseValue(0.0, 0.1))
# Left arm
x_bounds[-1]["min"][cls.YrotLeftUpperArm, 0] = -max_angle
x_bounds[-1]["max"][cls.YrotLeftUpperArm, 0] = 0
x_bounds[-2]["min"][cls.YrotLeftUpperArm, 2] = -max_angle
x_bounds[-2]["max"][cls.YrotLeftUpperArm, 2] = 0
define_loose_bounds(x_bounds[-1], cls.YrotLeftUpperArm, 2, LooseValue(-2.9, 0.1))
define_loose_bounds(x_bounds[-1], cls.ZrotLeftUpperArm, 2, LooseValue(0.0, 0.1))
@classmethod
def get_q_bounds(cls, half_twists: list, prefer_left: bool) -> dict:
nb_somersaults = len(half_twists)
is_forward = sum(half_twists) % 2 != 0
x_bounds = [
{
"min": cls.q_min_bounds.copy(),
"max": cls.q_max_bounds.copy(),
}
for _ in range(nb_somersaults + 1) # + 1 for landing
]
cls._fill_init_phase(x_bounds)
for phase in range(nb_somersaults):
cls._fill_somersault_phase(x_bounds, phase, half_twists)
cls._fill_landing_phase(x_bounds, half_twists)
if not is_forward:
invert_min_max(x_bounds, cls.Xrot)
if not prefer_left:
invert_min_max(x_bounds, cls.Zrot)
return x_bounds
@classmethod
def get_q_init(cls, half_twists: list = [], prefer_left: bool = True, q_bounds: np.ndarray = None) -> list:
x_bounds = q_bounds or cls.get_q_bounds(half_twists, prefer_left)
nb_phases = len(x_bounds)
x_inits = np.zeros((nb_phases, 2, cls.nb_q))
for phase in range(nb_phases):
x_inits[phase, 0] = (x_bounds[phase]["min"][:, 0] + x_bounds[phase]["max"][:, 0]) / 2
x_inits[phase, 1] = (x_bounds[phase]["min"][:, 2] + x_bounds[phase]["max"][:, 2]) / 2
return x_inits
@classmethod
def _fill_qdot_initial(cls, x_bounds: np.ndarray, final_time: float) -> None:
vzinit = 9.81 / 2 * final_time # initial Z velocity to land at final_time
x_bounds[0]["min"][:, 0] = [0] * cls.nb_qdot
x_bounds[0]["max"][:, 0] = [0] * cls.nb_qdot
x_bounds[0]["min"][: cls.Z, 0] = -0.5
x_bounds[0]["max"][: cls.Z, 0] = 0.5
x_bounds[0]["min"][cls.Z, 0] = vzinit - 5
x_bounds[0]["max"][cls.Z, 0] = vzinit + 5
x_bounds[0]["min"][cls.Xrot, 0] = 0.5
x_bounds[0]["max"][cls.Xrot, 0] = 50.0
@classmethod
def _fill_qdot_intermediary(cls, x_bounds: np.ndarray, final_time: float) -> None:
"""
vzinit = -g / 2 * final_time
vz(t) = -gt + vzinit
= g(final_time / 2 - t)
for any t, 0 <= t <= final_time
- vzinit <= vz(t) <= vzinit
we can use +- vzinit as Z velocity bounds (+- 10 to be sure)
"""
vzinit = 9.81 / 2 * final_time
nb_phases = len(x_bounds)
for phase in range(nb_phases):
if phase != 0:
# initial bounds, same as final bounds of previous phase
x_bounds[phase]["min"][:, 0] = x_bounds[phase - 1]["min"][:, 2]
x_bounds[phase]["max"][:, 0] = x_bounds[phase - 1]["max"][:, 2]
x_bounds[phase]["min"][:, 1:] = -100
x_bounds[phase]["max"][:, 1:] = 100
x_bounds[phase]["min"][: cls.Z, 1:] = -10
x_bounds[phase]["max"][: cls.Z, 1:] = 10
x_bounds[phase]["min"][cls.Z, 1:] = -vzinit - 10
x_bounds[phase]["max"][cls.Z, 1:] = vzinit + 10
x_bounds[phase]["min"][cls.Xrot, 1:] = 0.5
x_bounds[phase]["max"][cls.Xrot, 1:] = 50.0
@classmethod
def get_qdot_bounds(cls, nb_phases: int, final_time: float, is_forward: bool) -> dict:
x_bounds = [
{
"min": np.zeros((cls.nb_qdot, 3)) + cls.qdot_min,
"max": np.zeros((cls.nb_qdot, 3)) + cls.qdot_max,
}
for _ in range(nb_phases)
]
# Initial bounds
cls._fill_qdot_initial(x_bounds, final_time)
cls._fill_qdot_intermediary(x_bounds, final_time)
if not is_forward:
invert_min_max(x_bounds, cls.Xrot)
return x_bounds
@classmethod
def get_qdot_init(
cls,
nb_somersaults: int,
phase_durations: list[float],
is_forward: bool,
nb_phases: int,
) -> np.ndarray:
"""
Parameters
----------
nb_somersaults: int
The number of somersaults
phase_durations: list[float]
The duration of each phase
is_forward: bool
If the rotation is forward or backward
nb_phases: int
The number of phases
Returns
-------
The initial guess for the qdot
"""
g = 9.81
forward_factor = 1 if is_forward else -1
final_time = sum(phase_durations)
vzinit = g / 2 * final_time
qdot_init = np.zeros((nb_phases, cls.nb_qdot, 1))
qdot_init[:, cls.Xrot] = 2.5 * np.pi * nb_somersaults * forward_factor
for phase in range(nb_phases):
qdot_init[phase][cls.Z] = vzinit
return qdot_init
@classmethod
def get_tau_bounds(cls, nb_phases: int) -> list[dict[str, np.ndarray]]:
return [
{
"min": np.zeros((cls.nb_tau, 1)) + cls.tau_min,
"max": np.zeros((cls.nb_tau, 1)) + cls.tau_max,
}
for _ in range(nb_phases)
]
@classmethod
def get_tau_init(cls, nb_phases: int) -> np.ndarray:
return np.zeros((nb_phases, cls.nb_tau, 1)) + cls.tau_init