Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[RTR] some little test on preparing the data for animation of tracked markers. #714

Merged
merged 7 commits into from
Jul 13, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ def prepare_ocp(biorbd_model_path, phase_time, n_shooting, min_bound, max_bound)
tau_min, tau_max, tau_init = -500.0, 500.0, 0.0
activation_min, activation_max, activation_init = 0.0, 1.0, 0.5
dof_mapping = BiMappingList()
dof_mapping.add("tau", [None, None, None, 0], [3])
dof_mapping.add("tau", bimapping=None, to_second=[None, None, None, 0], to_first=[3])

# Add objective functions
objective_functions = ObjectiveList()
Expand Down Expand Up @@ -66,34 +66,36 @@ def prepare_ocp(biorbd_model_path, phase_time, n_shooting, min_bound, max_bound)

# Initialize x_bounds
x_bounds = BoundsList()
x_bounds.add(bounds=bio_model.bounds_from_ranges(["q", "qdot"]))
x_bounds[0][:, 0] = pose_at_first_node + [0] * n_qdot
x_bounds["q"] = bio_model.bounds_from_ranges("q")
x_bounds["q"][:, 0] = pose_at_first_node
x_bounds["qdot"] = bio_model.bounds_from_ranges("qdot")
x_bounds["qdot"][:, 0] = 0

# Initial guess
x_init = InitialGuessList()
x_init.add(pose_at_first_node + [0] * n_qdot)
x_init["q"] = pose_at_first_node
x_init["qdot"] = np.zeros((n_qdot,))

# Define control path constraint
u_bounds = BoundsList()
u_bounds.add(
[tau_min] * len(dof_mapping["tau"].to_first) + [activation_min] * bio_model.nb_muscles,
[tau_max] * len(dof_mapping["tau"].to_first) + [activation_max] * bio_model.nb_muscles,
)
u_bounds["tau"] = [tau_min] * len(dof_mapping["tau"].to_first), [tau_max] * len(dof_mapping["tau"].to_first)
u_bounds["muscles"] = [activation_min] * bio_model.nb_muscles, [activation_max] * bio_model.nb_muscles

u_init = InitialGuessList()
u_init.add([tau_init] * len(dof_mapping["tau"].to_first) + [activation_init] * bio_model.nb_muscles)
u_init["tau"] = [tau_init] * len(dof_mapping["tau"].to_first)
u_init["muscles"] = [activation_init] * bio_model.nb_muscles
# ------------- #

return OptimalControlProgram(
bio_model,
dynamics,
n_shooting,
phase_time,
x_init,
u_init,
x_bounds,
u_bounds,
objective_functions,
x_bounds=x_bounds,
u_bounds=u_bounds,
x_init=x_init,
u_init=u_init,
objective_functions=objective_functions,
constraints=constraints,
variable_mappings=dof_mapping,
assume_phase_dynamics=True,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@
SelectionMapping,
Dependency,
BoundsList,
Bounds,
InitialGuessList,
OdeSolver,
Solver,
Expand All @@ -38,12 +37,9 @@ def prepare_ocp(biorbd_model_path, phase_time, n_shooting, min_bound, ode_solver
# adds a bimapping to bimappinglist
# dof_mapping.add("tau", [None, None, None, 0], [3])
# easier way is to use SelectionMapping which is a subclass of biMapping
bimap = SelectionMapping(
nb_elements=bio_model.nb_dof,
independent_indices=(3,),
dependencies=(Dependency(dependent_index=None, reference_index=None, factor=None),),
)
dof_mapping.add(name="tau", bimapping=bimap)
dof_mapping = BiMappingList()
dof_mapping.add("tau", bimapping=None, to_second=[None, None, None, 0], to_first=[3])

# Add objective functions
objective_functions = ObjectiveList()
objective_functions.add(ObjectiveFcn.Mayer.MINIMIZE_PREDICTED_COM_HEIGHT, weight=-1)
Expand Down Expand Up @@ -77,34 +73,38 @@ def prepare_ocp(biorbd_model_path, phase_time, n_shooting, min_bound, ode_solver

# Initialize x_bounds
x_bounds = BoundsList()
x_bounds.add(bounds=bio_model.bounds_from_ranges(["q", "qdot"]))
x_bounds[0].concatenate(Bounds([activation_min] * n_mus, [activation_max] * n_mus))
x_bounds[0][:, 0] = pose_at_first_node + [0] * n_qdot + [0.5] * n_mus
x_bounds["q"] = bio_model.bounds_from_ranges("q")
x_bounds["q"][:, 0] = pose_at_first_node
x_bounds["qdot"] = bio_model.bounds_from_ranges("qdot")
x_bounds["qdot"][:, 0] = np.zeros((n_qdot,))
x_bounds["muscles"] = [[activation_min] * n_mus, [activation_max] * n_mus]
x_bounds["muscles"][:, 0] = np.zeros((n_mus,))

# Initial guess
x_init = InitialGuessList()
x_init.add(pose_at_first_node + [0] * n_qdot + [0.5] * n_mus)
x_init["q"] = pose_at_first_node
x_init["qdot"] = np.zeros((n_qdot,))
x_init["muscles"] = np.zeros((n_mus,))

# Define control path constraint
u_bounds = BoundsList()
u_bounds.add(
[torque_min] * len(dof_mapping["tau"].to_first) + [activation_min] * bio_model.nb_muscles,
[torque_max] * len(dof_mapping["tau"].to_first) + [activation_max] * bio_model.nb_muscles,
)
u_bounds["tau"] = [torque_min] * len(dof_mapping["tau"].to_first), [torque_max] * len(dof_mapping["tau"].to_first)
u_bounds["muscles"] = [activation_min] * bio_model.nb_muscles, [activation_max] * bio_model.nb_muscles

u_init = InitialGuessList()
u_init.add([torque_init] * len(dof_mapping["tau"].to_first) + [activation_init] * bio_model.nb_muscles)
u_init["tau"] = [torque_init] * len(dof_mapping["tau"].to_first)
u_init["muscles"] = [activation_init] * bio_model.nb_muscles
# ------------- #

return OptimalControlProgram(
bio_model,
dynamics,
n_shooting,
phase_time,
x_init,
u_init,
x_bounds,
u_bounds,
x_bounds=x_bounds,
u_bounds=u_bounds,
x_init=x_init,
u_init=u_init,
objective_functions=objective_functions,
constraints=constraints,
variable_mappings=dof_mapping,
Expand Down
26 changes: 25 additions & 1 deletion bioptim/interfaces/biomodel.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from typing import Protocol, Callable
import numpy as np
from typing import Protocol, Callable, Any

from casadi import MX, SX
from ..misc.mapping import BiMapping, BiMappingList
Expand Down Expand Up @@ -303,3 +304,26 @@ def partitioned_forward_dynamics(
ROBOTRAN: a powerful symbolic gnerator of multibody models, Mech. Sci., 4, 199–219,
https://doi.org/10.5194/ms-4-199-2013, 2013.
"""

@staticmethod
def animate(
solution: Any, show_now: bool = True, tracked_markers: list[np.ndarray, ...] = None, **kwargs: Any
) -> None | list:
"""
Animate a solution

Parameters
----------
solution: Any
The solution to animate
show_now: bool
If the animation should be shown immediately or not
tracked_markers: list[np.ndarray, ...]
The tracked markers (3, n_markers, n_frames)
kwargs: dict
The options to pass to the animator

Returns
-------
The animator object or None if show_now
"""
52 changes: 52 additions & 0 deletions bioptim/interfaces/biorbd_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
GeneralizedAcceleration,
)
from casadi import SX, MX, vertcat, horzcat, norm_fro
import numpy as np

from ..misc.utils import check_version
from ..limits.path_conditions import Bounds
Expand Down Expand Up @@ -512,6 +513,53 @@ def lagrangian(self, q: MX | SX, qdot: MX | SX) -> MX | SX:
qdot_biorbd = GeneralizedVelocity(qdot)
return self.model.Lagrangian(q_biorbd, qdot_biorbd).to_mx()

@staticmethod
def animate(
solution: Any, show_now: bool = True, tracked_markers: list[np.ndarray, ...] = None, **kwargs: Any
) -> None | list:
try:
import bioviz
except ModuleNotFoundError:
raise RuntimeError("bioviz must be install to animate the model")

check_version(bioviz, "2.3.0", "2.4.0")

states = solution.states
if not isinstance(states, (list, tuple)):
states = [states]

if tracked_markers is None:
tracked_markers = [None] * len(states)

all_bioviz = []
for idx_phase, data in enumerate(states):
if not isinstance(solution.ocp.nlp[idx_phase].model, BiorbdModel):
raise NotImplementedError("Animation is only implemented for biorbd models")

# This calls each of the function that modify the internal dynamic model based on the parameters
nlp = solution.ocp.nlp[idx_phase]

# noinspection PyTypeChecker
biorbd_model: BiorbdModel = nlp.model

all_bioviz.append(bioviz.Viz(biorbd_model.path, **kwargs))
all_bioviz[-1].load_movement(solution.ocp.nlp[idx_phase].variable_mappings["q"].to_second.map(data["q"]))

if tracked_markers[idx_phase] is not None:
all_bioviz[-1].load_experimental_markers(tracked_markers[idx_phase])

if show_now:
b_is_visible = [True] * len(all_bioviz)
while sum(b_is_visible):
for i, b in enumerate(all_bioviz):
if b.vtk_window.is_active:
b.update()
else:
b_is_visible[i] = False
return None
else:
return all_bioviz


class MultiBiorbdModel:
"""
Expand Down Expand Up @@ -1090,3 +1138,7 @@ def _qddot_mapping(self, mapping: BiMapping = None) -> dict:

def lagrangian(self):
raise NotImplementedError("lagrangian is not implemented yet for MultiBiorbdModel")

@staticmethod
def animate(solution: Any, show_now: bool = True, tracked_markers: list = [], **kwargs: Any) -> None | list:
raise NotImplementedError("animate is not implemented yet for MultiBiorbdModel")
106 changes: 68 additions & 38 deletions bioptim/optimization/solution.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@

import numpy as np
from scipy import interpolate as sci_interp
from scipy.interpolate import interp1d
from casadi import vertcat, DM, Function
from matplotlib import pyplot as plt

Expand All @@ -24,6 +25,7 @@
from ..optimization.optimization_vector import OptimizationVectorHelper
from ..dynamics.ode_solver import OdeSolver
from ..interfaces.solve_ivp_interface import solve_ivp_interface, solve_ivp_bioptim_interface
from ..interfaces.biorbd_model import BiorbdModel


class Solution:
Expand Down Expand Up @@ -1350,7 +1352,12 @@ def graphs(
plt.show()

def animate(
self, n_frames: int = 0, shooting_type: Shooting = None, show_now: bool = True, **kwargs: Any
self,
n_frames: int = 0,
shooting_type: Shooting = None,
show_now: bool = True,
show_tracked_markers: bool = False,
**kwargs: Any,
) -> None | list:
"""
Animate the simulation
Expand All @@ -1364,6 +1371,8 @@ def animate(
The Shooting type to animate
show_now: bool
If the bioviz exec() function should be called automatically. This is blocking method
show_tracked_markers: bool
If the tracked markers should be displayed
kwargs: Any
Any parameters to pass to bioviz

Expand All @@ -1372,61 +1381,82 @@ def animate(
A list of bioviz structures (one for each phase). So one can call exec() by hand
"""

try:
import bioviz
except ModuleNotFoundError:
raise RuntimeError("bioviz must be install to animate the model")

from ..interfaces.biorbd_model import BiorbdModel

check_version(bioviz, "2.3.0", "2.4.0")

data_to_animate = self.integrate(shooting_type=shooting_type) if shooting_type else self.copy()
if n_frames == 0:
try:
data_to_animate = data_to_animate.interpolate(sum(self.ns))
data_to_animate = data_to_animate.interpolate(sum(self.ns) + 1)
except RuntimeError:
pass

elif n_frames > 0:
data_to_animate = data_to_animate.interpolate(n_frames)

states = data_to_animate.states
if not isinstance(states, (list, tuple)):
states = [states]
if show_tracked_markers and len(self.ocp.nlp) == 1:
tracked_markers = self._prepare_tracked_markers_for_animation(n_shooting=n_frames)
elif show_tracked_markers and len(self.ocp.nlp) > 1:
raise NotImplementedError(
"Tracking markers is not implemented for multiple phases. "
"Set show_tracked_markers to False such that sol.animate(show_tracked_markers=False)."
)
else:
tracked_markers = [None for _ in range(len(self.ocp.nlp))]

# assuming that all the models or the same type.
self._check_models_comes_from_same_super_class()
self.ocp.nlp[0].model.animate(
solution=data_to_animate,
show_now=show_now,
tracked_markers=tracked_markers,
**kwargs,
)

def _check_models_comes_from_same_super_class(self):
"""Check that all the models comes from the same super class"""
for i, nlp in enumerate(self.ocp.nlp):
model_super_classes = nlp.model.__class__.mro()[:-1] # remove object class
nlps = self.ocp.nlp.copy()
del nlps[i]
for j, sub_nlp in enumerate(nlps):
if not any([isinstance(sub_nlp.model, super_class) for super_class in model_super_classes]):
raise RuntimeError(
f"The animation is only available for compatible models. "
f"Here, the model of phase {i} is of type {nlp.model.__class__.__name__} and the model of "
f"phase {j + 1 if i < j else j} is of type {sub_nlp.model.__class__.__name__} and "
f"they don't share the same super class."
)

all_bioviz = []
for idx_phase, data in enumerate(states):
if not isinstance(self.ocp.nlp[idx_phase].model, BiorbdModel):
raise NotImplementedError("Animation is only implemented for biorbd models")
def _prepare_tracked_markers_for_animation(self, n_shooting: int = None) -> list:
"""Prepare the markers which are tracked to the animation"""

# This calls each of the function that modify the internal dynamic model based on the parameters
nlp = self.ocp.nlp[idx_phase]
n_frames = sum(self.ns) + 1 if n_shooting is None else n_shooting + 1

# noinspection PyTypeChecker
biorbd_model: BiorbdModel = nlp.model
all_tracked_markers = []

all_bioviz.append(bioviz.Viz(biorbd_model.path, **kwargs))
all_bioviz[-1].load_movement(self.ocp.nlp[idx_phase].variable_mappings["q"].to_second.map(data["q"]))
for objective in self.ocp.nlp[idx_phase].J:
for phase, nlp in enumerate(self.ocp.nlp):
tracked_markers = None
for objective in nlp.J:
if objective.target is not None:
if objective.type in (
ObjectiveFcn.Mayer.TRACK_MARKERS,
ObjectiveFcn.Lagrange.TRACK_MARKERS,
) and objective.node[0] in (Node.ALL, Node.ALL_SHOOTING):
all_bioviz[-1].load_experimental_markers(objective.target[0])

if show_now:
b_is_visible = [True] * len(all_bioviz)
while sum(b_is_visible):
for i, b in enumerate(all_bioviz):
if b.vtk_window.is_active:
b.update()
else:
b_is_visible[i] = False
return None
else:
return all_bioviz
tracked_markers = np.full((3, nlp.model.nb_markers, self.ns[phase] + 1), np.nan)
for i in range(len(objective.rows)):
tracked_markers[objective.rows[i], objective.cols, :] = objective.target[0][i, :, :]
missing_row = np.where(np.isnan(tracked_markers))[0]
if missing_row.size > 0:
tracked_markers[missing_row, :, :] = 0

# interpolation
if n_frames > 0 and tracked_markers is not None:
x = np.linspace(0, self.ns[phase], self.ns[phase] + 1)
xnew = np.linspace(0, self.ns[phase], n_frames)
f = interp1d(x, tracked_markers, kind="cubic")
tracked_markers = f(xnew)

all_tracked_markers.append(tracked_markers)

return all_tracked_markers

def _get_penalty_cost(self, nlp, penalty):
phase_idx = nlp.phase_idx
Expand Down
Loading