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.
Jump to
Jump to file
Failed to load files.
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 @@
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 warning on line 523 in bioptim/interfaces/biorbd_model.py

View check run for this annotation

Codecov / codecov/patch

bioptim/interfaces/biorbd_model.py#L520-L523

Added lines #L520 - L523 were not covered by tests

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

Check warning on line 525 in bioptim/interfaces/biorbd_model.py

View check run for this annotation

Codecov / codecov/patch

bioptim/interfaces/biorbd_model.py#L525

Added line #L525 was not covered by tests

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

Check warning on line 529 in bioptim/interfaces/biorbd_model.py

View check run for this annotation

Codecov / codecov/patch

bioptim/interfaces/biorbd_model.py#L527-L529

Added lines #L527 - L529 were not covered by tests

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

Check warning on line 532 in bioptim/interfaces/biorbd_model.py

View check run for this annotation

Codecov / codecov/patch

bioptim/interfaces/biorbd_model.py#L531-L532

Added lines #L531 - L532 were not covered by tests

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")

Check warning on line 537 in bioptim/interfaces/biorbd_model.py

View check run for this annotation

Codecov / codecov/patch

bioptim/interfaces/biorbd_model.py#L534-L537

Added lines #L534 - L537 were not covered by tests

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

Check warning on line 540 in bioptim/interfaces/biorbd_model.py

View check run for this annotation

Codecov / codecov/patch

bioptim/interfaces/biorbd_model.py#L540

Added line #L540 was not covered by tests

# noinspection PyTypeChecker
biorbd_model: BiorbdModel = nlp.model

Check warning on line 543 in bioptim/interfaces/biorbd_model.py

View check run for this annotation

Codecov / codecov/patch

bioptim/interfaces/biorbd_model.py#L543

Added line #L543 was not covered by tests

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"]))

Check warning on line 546 in bioptim/interfaces/biorbd_model.py

View check run for this annotation

Codecov / codecov/patch

bioptim/interfaces/biorbd_model.py#L545-L546

Added lines #L545 - L546 were not covered by tests

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

Check warning on line 549 in bioptim/interfaces/biorbd_model.py

View check run for this annotation

Codecov / codecov/patch

bioptim/interfaces/biorbd_model.py#L548-L549

Added lines #L548 - L549 were not covered by tests

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()

Check warning on line 556 in bioptim/interfaces/biorbd_model.py

View check run for this annotation

Codecov / codecov/patch

bioptim/interfaces/biorbd_model.py#L551-L556

Added lines #L551 - L556 were not covered by tests
else:
b_is_visible[i] = False
return None

Check warning on line 559 in bioptim/interfaces/biorbd_model.py

View check run for this annotation

Codecov / codecov/patch

bioptim/interfaces/biorbd_model.py#L558-L559

Added lines #L558 - L559 were not covered by tests
else:
return all_bioviz

Check warning on line 561 in bioptim/interfaces/biorbd_model.py

View check run for this annotation

Codecov / codecov/patch

bioptim/interfaces/biorbd_model.py#L561

Added line #L561 was not covered by tests


class MultiBiorbdModel:
"""
Expand Down Expand Up @@ -1090,3 +1138,7 @@

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")

Check warning on line 1144 in bioptim/interfaces/biorbd_model.py

View check run for this annotation

Codecov / codecov/patch

bioptim/interfaces/biorbd_model.py#L1144

Added line #L1144 was not covered by tests
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 @@
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 @@
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 @@
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)

Check warning on line 1387 in bioptim/optimization/solution.py

View check run for this annotation

Codecov / codecov/patch

bioptim/optimization/solution.py#L1387

Added line #L1387 was not covered by tests
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(

Check warning on line 1397 in bioptim/optimization/solution.py

View check run for this annotation

Codecov / codecov/patch

bioptim/optimization/solution.py#L1394-L1397

Added lines #L1394 - L1397 were not covered by tests
"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))]

Check warning on line 1402 in bioptim/optimization/solution.py

View check run for this annotation

Codecov / codecov/patch

bioptim/optimization/solution.py#L1402

Added line #L1402 was not covered by tests

# assuming that all the models or the same type.
self._check_models_comes_from_same_super_class()
self.ocp.nlp[0].model.animate(

Check warning on line 1406 in bioptim/optimization/solution.py

View check run for this annotation

Codecov / codecov/patch

bioptim/optimization/solution.py#L1405-L1406

Added lines #L1405 - L1406 were not covered by tests
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