diff --git a/README.md b/README.md index a4d3086d..be42fc32 100644 --- a/README.md +++ b/README.md @@ -69,7 +69,7 @@ conda install -cconda-forge conda-libmamba-solver After, install the dependencies ```bash -conda install numpy matplotlib pytest casadi biorbd -cconda-forge --solver=libmamba +conda install numpy matplotlib pytest casadi biorbd pyorerun -cconda-forge --solver=libmamba ``` Finally, install the bioptim setup.py file located in your cocofest/external/bioptim folder diff --git a/cocofest/__init__.py b/cocofest/__init__.py index bc968cc5..483a2870 100644 --- a/cocofest/__init__.py +++ b/cocofest/__init__.py @@ -26,3 +26,4 @@ ) from .result.plot import PlotCyclingResult from .result.pickle import SolutionToPickle +from .result.animate import PickleAnimate diff --git a/cocofest/custom_objectives.py b/cocofest/custom_objectives.py index a849fd5c..79d7c43f 100644 --- a/cocofest/custom_objectives.py +++ b/cocofest/custom_objectives.py @@ -88,7 +88,7 @@ def minimize_overall_muscle_fatigue(controller: PenaltyController) -> MX: muscle_fatigue = horzcat( *[controller.states["A_" + muscle_name_list[x]].cx for x in range(len(muscle_name_list))] ) - return 1 / sum1(muscle_fatigue) + return sum1(muscle_fatigue) @staticmethod def minimize_overall_muscle_force_production(controller: PenaltyController) -> MX: @@ -106,6 +106,6 @@ def minimize_overall_muscle_force_production(controller: PenaltyController) -> M """ muscle_name_list = controller.model.bio_model.muscle_names muscle_force = horzcat( - *[controller.states["F_" + muscle_name_list[x]].cx ** 3 for x in range(len(muscle_name_list))] + *[controller.states["F_" + muscle_name_list[x]].cx for x in range(len(muscle_name_list))] ) return sum1(muscle_force) diff --git a/cocofest/examples/dynamics/reaching_task/reaching_task_pulse_duration_optimization.py b/cocofest/examples/dynamics/reaching_task/reaching_task_pulse_duration_optimization.py index 765b0bf4..67e7aa9d 100644 --- a/cocofest/examples/dynamics/reaching_task/reaching_task_pulse_duration_optimization.py +++ b/cocofest/examples/dynamics/reaching_task/reaching_task_pulse_duration_optimization.py @@ -5,19 +5,17 @@ The files will contain the time, states, controls and parameters of the ocp. """ -import pickle - from bioptim import ( Axis, ConstraintFcn, ConstraintList, Node, Solver, - SolutionMerge, ) -from cocofest import DingModelPulseDurationFrequencyWithFatigue, OcpFesMsk +from cocofest import DingModelPulseDurationFrequencyWithFatigue, OcpFesMsk, SolutionToPickle +# Scaling alpha_a and a_scale parameters for each muscle proportionally to the muscle PCSA and fiber type 2 proportion # Fiber type proportion from [1] biceps_fiber_type_2_proportion = 0.607 triceps_fiber_type_2_proportion = 0.465 @@ -32,13 +30,12 @@ ] # PCSA (cm²) from [2] -biceps_pcsa = 12.7 triceps_pcsa = 28.3 +biceps_pcsa = 12.7 brachioradialis_pcsa = 11.6 - -biceps_a_scale_proportion = 12.7 / 28.3 triceps_a_scale_proportion = 1 -brachioradialis_a_scale_proportion = 11.6 / 28.3 +biceps_a_scale_proportion = biceps_pcsa / triceps_pcsa +brachioradialis_a_scale_proportion = brachioradialis_pcsa / triceps_pcsa a_scale_proportion_list = [ biceps_a_scale_proportion, biceps_a_scale_proportion, @@ -48,6 +45,8 @@ brachioradialis_a_scale_proportion, ] +# Build the functional electrical stimulation models according +# to number and name of muscle in the musculoskeletal model used fes_muscle_models = [ DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), DingModelPulseDurationFrequencyWithFatigue(muscle_name="BICshort"), @@ -57,6 +56,7 @@ DingModelPulseDurationFrequencyWithFatigue(muscle_name="BRA"), ] +# Applying the scaling for i in range(len(fes_muscle_models)): fes_muscle_models[i].alpha_a = fes_muscle_models[i].alpha_a * alpha_a_proportion_list[i] fes_muscle_models[i].a_scale = fes_muscle_models[i].a_scale * a_scale_proportion_list[i] @@ -64,7 +64,8 @@ minimum_pulse_duration = DingModelPulseDurationFrequencyWithFatigue().pd0 pickle_file_list = ["minimize_muscle_fatigue.pkl", "minimize_muscle_force.pkl"] n_stim = 60 -n_shooting = 2 +n_shooting = 25 +# Step time of 1ms -> 1sec / (40Hz * 25) = 0.001s constraint = ConstraintList() constraint.add( @@ -100,22 +101,7 @@ ) sol = ocp.solve(Solver.IPOPT(_max_iter=10000)) - - time = sol.decision_time(to_merge=[SolutionMerge.PHASES, SolutionMerge.NODES]) - states = sol.decision_states(to_merge=[SolutionMerge.PHASES, SolutionMerge.NODES]) - controls = sol.decision_controls(to_merge=[SolutionMerge.PHASES, SolutionMerge.NODES]) - parameters = sol.decision_parameters() - - dictionary = { - "time": time, - "states": states, - "controls": controls, - "parameters": parameters, - } - - with open("result_file/pulse_duration_" + pickle_file_list[i], "wb") as file: - pickle.dump(dictionary, file) - + SolutionToPickle(sol, "pulse_duration_" + pickle_file_list[i], "result_file/").pickle() # [1] Dahmane, R., Djordjevič, S., Šimunič, B., & Valenčič, V. (2005). # Spatial fiber type distribution in normal human muscle: histochemical and tensiomyographical evaluation. diff --git a/cocofest/examples/dynamics/reaching_task/result_file/pulse_duration_minimize_muscle_fatigue.pkl b/cocofest/examples/dynamics/reaching_task/result_file/pulse_duration_minimize_muscle_fatigue.pkl index a3b62e9b..c30b4a70 100644 Binary files a/cocofest/examples/dynamics/reaching_task/result_file/pulse_duration_minimize_muscle_fatigue.pkl and b/cocofest/examples/dynamics/reaching_task/result_file/pulse_duration_minimize_muscle_fatigue.pkl differ diff --git a/cocofest/examples/dynamics/reaching_task/result_file/pulse_duration_minimize_muscle_force.pkl b/cocofest/examples/dynamics/reaching_task/result_file/pulse_duration_minimize_muscle_force.pkl index 7bcbf5e6..3037308a 100644 Binary files a/cocofest/examples/dynamics/reaching_task/result_file/pulse_duration_minimize_muscle_force.pkl and b/cocofest/examples/dynamics/reaching_task/result_file/pulse_duration_minimize_muscle_force.pkl differ diff --git a/cocofest/examples/dynamics/reaching_task/visualization/animation.py b/cocofest/examples/dynamics/reaching_task/visualization/animation.py new file mode 100644 index 00000000..bcc8e84d --- /dev/null +++ b/cocofest/examples/dynamics/reaching_task/visualization/animation.py @@ -0,0 +1,8 @@ +from cocofest import PickleAnimate + +PickleAnimate("../result_file/pulse_duration_minimize_muscle_fatigue.pkl").animate( + model_path="../../../msk_models/arm26.bioMod" +) +PickleAnimate("../result_file/pulse_duration_minimize_muscle_fatigue.pkl").multiple_animations( + ["../result_file/pulse_duration_minimize_muscle_force.pkl"], model_path="../../../msk_models/arm26.bioMod" +) diff --git a/cocofest/examples/dynamics/reaching_task/visualization/force_and_fatigue_subplot.py b/cocofest/examples/dynamics/reaching_task/visualization/force_and_fatigue_subplot.py new file mode 100644 index 00000000..cabea30b --- /dev/null +++ b/cocofest/examples/dynamics/reaching_task/visualization/force_and_fatigue_subplot.py @@ -0,0 +1,224 @@ +""" +This script is used to make the graph of the muscle force and fatigue for the reaching task. +The data used to make the graph is from the result file of the optimization. +""" + +import pickle +import matplotlib.pyplot as plt +import numpy as np + +pickle_path = [ + r"../result_file/pulse_duration_minimize_muscle_force.pkl", + r"../result_file/pulse_duration_minimize_muscle_fatigue.pkl", +] + +with open(pickle_path[0], "rb") as f: + data_minimize_force = pickle.load(f) + +with open(pickle_path[1], "rb") as f: + data_minimize_fatigue = pickle.load(f) + +force_muscle_keys = ["F_BIClong", "F_BICshort", "F_TRIlong", "F_TRIlat", "F_TRImed", "F_BRA"] +fatigue_muscle_keys = ["A_BIClong", "A_BICshort", "A_TRIlong", "A_TRIlat", "A_TRImed", "A_BRA"] +muscle_names = ["BIClong", "BICshort", "TRIlong", "TRIlat", "TRImed", "BRA"] + +# Force graph +fig, axs = plt.subplots(3, 2, figsize=(6, 7)) +fig.suptitle("Muscle force", fontsize=20, fontweight="bold", fontname="Times New Roman") +index = 0 + +for j in range(2): + for i in range(3): + axs[i][j].set_xlim(left=0, right=1.5) + axs[i][j].set_ylim(bottom=0, top=250) + if j == 0: + plt.setp( + axs[i][j], + xticks=[0, 0.5, 1, 1.5], + xticklabels=[], + yticks=[0, 75, 150, 225], + yticklabels=[0, 75, 150, 225], + ) + + else: + plt.setp( + axs[i][j], + xticks=[0, 0.5, 1, 1.5], + xticklabels=[], + yticks=[0, 75, 150, 225], + yticklabels=[], + ) + + if i == 2: + plt.setp( + axs[i][j], + xticks=[0, 0.5, 1, 1.5], + xticklabels=[0, 0.5, 1, 1.5], + ) + + axs[i][j].plot(data_minimize_force["time"], data_minimize_force["states"][force_muscle_keys[index]], lw=5) + axs[i][j].plot(data_minimize_fatigue["time"], data_minimize_fatigue["states"][force_muscle_keys[index]], lw=5) + axs[i][j].text( + 0.5, + 0.9, + f"{muscle_names[index]}", + transform=axs[i][j].transAxes, + ha="center", + va="center", + fontsize=14, + weight="bold", + font="Times New Roman", + ) + + labels = axs[i][j].get_xticklabels() + axs[i][j].get_yticklabels() + [label.set_fontname("Times New Roman") for label in labels] + [label.set_fontsize(14) for label in labels] + + index += 1 + +fig.text(0.5, 0.02, "Time (s)", ha="center", va="center", fontsize=18, weight="bold", font="Times New Roman") +fig.text( + 0.025, + 0.5, + "Force (N)", + ha="center", + va="center", + rotation="vertical", + fontsize=18, + weight="bold", + font="Times New Roman", +) +fig.legend( + ["Force", "Fatigue"], loc="upper right", ncol=1, prop={"family": "Times New Roman", "size": 14, "weight": "bold"} +) +plt.show() + +# Joint angle graph +joint_keys = ["Shoulder", "Elbow"] +fig, axs = plt.subplots(2, 1, figsize=(3, (2 / 3) * 7)) +fig.suptitle("Joint angle", fontsize=20, fontweight="bold", fontname="Times New Roman") + +for i in range(2): + axs[i].set_xlim(left=0, right=1.5) + + if i == 1: + plt.setp( + axs[i], + xticks=[0, 0.5, 1, 1.5], + xticklabels=[0, 0.5, 1, 1.5], + ) + else: + plt.setp( + axs[i], + xticks=[0, 0.5, 1, 1.5], + xticklabels=[], + ) + + force_angles = data_minimize_force["states"]["q"][i] * 180 / 3.14 + fatigue_angles = data_minimize_fatigue["states"]["q"][i] * 180 / 3.14 + + axs[i].plot(data_minimize_force["time"], force_angles, lw=5) + axs[i].plot(data_minimize_fatigue["time"], fatigue_angles, lw=5) + axs[i].text( + 0.05, + 0.9, + f"{joint_keys[i]}", + transform=axs[i].transAxes, + ha="left", + va="center", + fontsize=14, + weight="bold", + font="Times New Roman", + ) + labels = axs[i].get_xticklabels() + axs[i].get_yticklabels() + [label.set_fontname("Times New Roman") for label in labels] + [label.set_fontsize(14) for label in labels] + +fig.text( + 0.05, + 0.5, + "Joint angle (°)", + ha="center", + va="center", + rotation="vertical", + fontsize=18, + weight="bold", + font="Times New Roman", +) +axs[1].text(0.75, -25, "Time (s)", ha="center", va="center", fontsize=18, weight="bold", font="Times New Roman") +fig.legend( + ["Force", "Fatigue"], loc="upper right", ncol=1, prop={"family": "Times New Roman", "size": 14, "weight": "bold"} +) +plt.show() + +# Fatigue graph +a_list = ["A_BIClong", "A_BICshort", "A_TRIlong", "A_TRIlat", "A_TRImed", "A_BRA"] +a_sum_base_line = 0 +a_force_sum_list = [] +a_fatigue_sum_list = [] +for key_a in a_list: + a_sum_base_line += data_minimize_force["states"][key_a][0] +for i in range(len(data_minimize_force["time"])): + a_force_sum = 0 + a_fatigue_sum = 0 + for key_a in a_list: + a_force_sum += data_minimize_force["states"][key_a][i] + a_fatigue_sum += data_minimize_fatigue["states"][key_a][i] + + a_force_sum_list.append(a_force_sum) + a_fatigue_sum_list.append(a_fatigue_sum) + +a_force_diff_list = [] +a_fatigue_diff_list = [] +fatigue_minimization_percentage_gain_list = [] +for i in range(len(data_minimize_force["time"])): + a_force_diff_list.append((a_force_sum_list[i] - a_force_sum_list[0]) * 1000) + a_fatigue_diff_list.append((a_fatigue_sum_list[i] - a_fatigue_sum_list[0]) * 1000) + + fatigue_minimization_percentage_gain_list.append( + (a_fatigue_sum_list[i] - a_force_sum_list[i]) / (a_force_sum_list[0] - a_force_sum_list[-1]) * 100 + ) + +fig, axs = plt.subplots(1, 1, figsize=(3, (1 / 3) * 7)) +fig.suptitle("Muscle fatigue", fontsize=20, fontweight="bold", fontname="Times New Roman") + +axs.set_xlim(left=0, right=1.5) +plt.setp( + axs, + xticks=[0, 0.5, 1, 1.5], + xticklabels=[0, 0.5, 1, 1.5], +) + +a_force_sum_percentage = (np.array(a_force_sum_list) / a_sum_base_line) * 100 +a_fatigue_sum_percentage = (np.array(a_fatigue_sum_list) / a_sum_base_line) * 100 + +axs.plot(data_minimize_force["time"], a_force_sum_percentage, lw=5) +axs.plot(data_minimize_force["time"], a_fatigue_sum_percentage, lw=5) + +axs.set_xlim(left=0, right=1.5) + +plt.setp( + axs, + xticks=[0, 0.5, 1, 1.5], + xticklabels=[0, 0.5, 1, 1.5], +) + +labels = axs.get_xticklabels() + axs.get_yticklabels() +[label.set_fontname("Times New Roman") for label in labels] +[label.set_fontsize(14) for label in labels] +fig.text( + 0.05, + 0.5, + "Fatigue percentage (%)", + ha="center", + va="center", + rotation="vertical", + fontsize=18, + weight="bold", + font="Times New Roman", +) +axs.text(0.75, 96.3, "Time (s)", ha="center", va="center", fontsize=18, weight="bold", font="Times New Roman") +plt.legend( + ["Force", "Fatigue"], loc="upper right", ncol=1, prop={"family": "Times New Roman", "size": 14, "weight": "bold"} +) +plt.show() diff --git a/cocofest/examples/dynamics/reaching_task/make_gaph.py b/cocofest/examples/dynamics/reaching_task/visualization/make_gaph.py similarity index 95% rename from cocofest/examples/dynamics/reaching_task/make_gaph.py rename to cocofest/examples/dynamics/reaching_task/visualization/make_gaph.py index 81c515e8..889cc6cb 100644 --- a/cocofest/examples/dynamics/reaching_task/make_gaph.py +++ b/cocofest/examples/dynamics/reaching_task/visualization/make_gaph.py @@ -10,8 +10,8 @@ chosen_graph_to_plot = "duration" duration_path = [ - r"result_file/pulse_duration_minimize_muscle_force.pkl", - r"result_file/pulse_duration_minimize_muscle_fatigue.pkl", + r"../result_file/pulse_duration_minimize_muscle_force.pkl", + r"../result_file/pulse_duration_minimize_muscle_fatigue.pkl", ] chosen_graph_to_plot_path = duration_path if chosen_graph_to_plot == "duration" else None @@ -74,14 +74,14 @@ if i == 0 and j == 0: axs[i][j].plot( data_minimize_force["time"], - data_minimize_force["states"][force_muscle_keys[index]][0], + data_minimize_force["states"][force_muscle_keys[index]], ms=4, linewidth=5.0, label="Minimizing force", ) axs[i][j].plot( data_minimize_fatigue["time"], - data_minimize_fatigue["states"][force_muscle_keys[index]][0], + data_minimize_fatigue["states"][force_muscle_keys[index]], ms=4, linewidth=5.0, label="Minimizing fatigue", @@ -89,13 +89,13 @@ else: axs[i][j].plot( data_minimize_force["time"], - data_minimize_force["states"][force_muscle_keys[index]][0], + data_minimize_force["states"][force_muscle_keys[index]], ms=4, linewidth=5.0, ) axs[i][j].plot( data_minimize_fatigue["time"], - data_minimize_fatigue["states"][force_muscle_keys[index]][0], + data_minimize_fatigue["states"][force_muscle_keys[index]], ms=4, linewidth=5.0, ) @@ -177,13 +177,13 @@ a_force_sum_list = [] a_fatigue_sum_list = [] for key_a in a_list: - a_sum_base_line += data_minimize_force["states"][key_a][0][0] + a_sum_base_line += data_minimize_force["states"][key_a][0] for i in range(len(data_minimize_force["time"])): a_force_sum = 0 a_fatigue_sum = 0 for key_a in a_list: - a_force_sum += data_minimize_force["states"][key_a][0][i] - a_fatigue_sum += data_minimize_fatigue["states"][key_a][0][i] + a_force_sum += data_minimize_force["states"][key_a][i] + a_fatigue_sum += data_minimize_fatigue["states"][key_a][i] a_force_sum_list.append(a_force_sum) a_fatigue_sum_list.append(a_fatigue_sum) diff --git a/cocofest/examples/dynamics/reaching_task/visualization/pulse_duration_subplot.py b/cocofest/examples/dynamics/reaching_task/visualization/pulse_duration_subplot.py new file mode 100644 index 00000000..b02df00b --- /dev/null +++ b/cocofest/examples/dynamics/reaching_task/visualization/pulse_duration_subplot.py @@ -0,0 +1,68 @@ +""" +This script is used to make the graph of the muscle force and fatigue for the reaching task. +The data used to make the graph is from the result file of the optimization. +The available graphs are: duration +""" + +import pickle +import matplotlib.pyplot as plt +import matplotlib.colors as colors +import matplotlib.cm as cmx + +pickle_path = [ + r"../result_file/pulse_duration_minimize_muscle_force.pkl", + r"../result_file/pulse_duration_minimize_muscle_fatigue.pkl", +] + +with open(pickle_path[0], "rb") as f: + data_minimize_force = pickle.load(f) + +with open(pickle_path[1], "rb") as f: + data_minimize_fatigue = pickle.load(f) + +pulse_duration_keys = list(data_minimize_fatigue["parameters"].keys()) +muscle_names = ["BIClong", "BICshort", "TRIlong", "TRIlat", "TRImed", "BRA"] +nb_stim = len(data_minimize_fatigue["parameters"][pulse_duration_keys[0]]) +width = round(data_minimize_fatigue["time"][-1], 2) / nb_stim + +pw_data_list = [data_minimize_force["parameters"], data_minimize_fatigue["parameters"]] + +pw_list = [] +for j in range(2): + pw_list.append([pw_data_list[j][pulse_duration_keys[i]] * 1000000 for i in range(len(pulse_duration_keys))]) + +plasma = cm = plt.get_cmap("plasma") +cNorm = colors.Normalize(vmin=100, vmax=600) +scalarMap = cmx.ScalarMappable(norm=cNorm, cmap=plasma) + + +def plot_graph(datas): + fig, axs = plt.subplots(6, 1, figsize=(5, 3), constrained_layout=True) + + for i in range(len(pulse_duration_keys)): + axs[i].set_xlim(left=0, right=1.5) + plt.setp( + axs[i], + xticks=[0, 0.2, 0.4, 0.6, 0.8, 1, 1.2, 1.4], + xticklabels=[], + ) + if i == len(pulse_duration_keys) - 1: + plt.setp( + axs[i], + xticks=[0, 0.2, 0.4, 0.6, 0.8, 1, 1.2, 1.4], + xticklabels=[0, 0.2, 0.4, 0.6, 0.8, 1, 1.2, 1.4], + ) + axs[i].set_xlabel("Time (s)") + + for j in range(nb_stim): + value = datas[i][j] + + color = scalarMap.to_rgba(value) + axs[i].barh(muscle_names[i], width, left=j * width, height=0.5, color=color) + + fig.colorbar(scalarMap, ax=axs, orientation="vertical", label="Pulse duration (us)") + plt.show() + + +for data in pw_list: + plot_graph(data) diff --git a/cocofest/fourier_approx.py b/cocofest/fourier_approx.py index 8196f614..e9a54d64 100644 --- a/cocofest/fourier_approx.py +++ b/cocofest/fourier_approx.py @@ -12,8 +12,8 @@ def __init__(self): def compute_real_fourier_coeffs(self, x, y, n): result = [] for i in range(n + 1): - an = (2.0 / self.p) * spi.trapz(y * np.cos(2 * np.pi * i * x / self.p), x) - bn = (2.0 / self.p) * spi.trapz(y * np.sin(2 * np.pi * i * x / self.p), x) + an = (2.0 / self.p) * spi.trapezoid(y * np.cos(2 * np.pi * i * x / self.p), x) + bn = (2.0 / self.p) * spi.trapezoid(y * np.sin(2 * np.pi * i * x / self.p), x) result.append((an, bn)) return np.array(result) diff --git a/cocofest/optimization/fes_ocp_dynamics.py b/cocofest/optimization/fes_ocp_dynamics.py index d2af0854..2913e298 100644 --- a/cocofest/optimization/fes_ocp_dynamics.py +++ b/cocofest/optimization/fes_ocp_dynamics.py @@ -753,26 +753,29 @@ def _set_objective( ) if minimize_muscle_fatigue: - for i in range(n_stim): - objective_functions.add( - CustomObjective.minimize_overall_muscle_fatigue, - custom_type=ObjectiveFcn.Mayer, - node=Node.ALL, - quadratic=True, - weight=1, - phase=i, - ) - - if minimize_muscle_force: + # for i in range(n_stim): objective_functions.add( - CustomObjective.minimize_overall_muscle_force_production, + CustomObjective.minimize_overall_muscle_fatigue, + # custom_type=ObjectiveFcn.Lagrange, custom_type=ObjectiveFcn.Mayer, node=Node.END, quadratic=True, - weight=1, + weight=-1, phase=n_stim - 1, ) + if minimize_muscle_force: + for i in range(n_stim): + objective_functions.add( + CustomObjective.minimize_overall_muscle_force_production, + custom_type=ObjectiveFcn.Lagrange, + # custom_type=ObjectiveFcn.Mayer, + # node=Node.ALL, + quadratic=True, + weight=1, + phase=i, + ) + if time_min and time_max: for i in range(n_stim): objective_functions.add( diff --git a/cocofest/result/animate.py b/cocofest/result/animate.py new file mode 100644 index 00000000..c8e80c68 --- /dev/null +++ b/cocofest/result/animate.py @@ -0,0 +1,72 @@ +import pyorerun as prr +import biorbd +import pickle +import numpy as np + + +class PickleAnimate: + def __init__(self, path: str): + self.path = path + self.data = None + self.model = None + self.time = None + self.state_q = None + self.frames = None + + def load(self): + with open(self.path, "rb") as f: + self.data = pickle.load(f) + + model_path = self.data["bio_model_path"] if "bio_model_path" in self.data.keys() else None + if model_path is None: + raise ValueError("The bio model path is not available, please provide it to animate the solution.") + + # Load a predefined model + if self.model is None: + self.model = biorbd.Model(model_path) + self.time = self.data["time"] + self.state_q = self.data["states"]["q"] + self.frames = self.state_q.shape[1] + + def animate(self, model_path: str = None): + if model_path: + self.model = biorbd.Model(model_path) + self.load() + + # pyorerun animation + prr_model = prr.BiorbdModel.from_biorbd_object(self.model) + + nb_seconds = self.time[-1] + t_span = np.linspace(0, nb_seconds, self.frames) + + viz = prr.PhaseRerun(t_span) + viz.add_animated_model(prr_model, self.state_q) + viz.rerun("msk_model") + + def multiple_animations(self, additional_path: list[str], model_path: str = None): + if model_path: + self.model = biorbd.Model(model_path) + self.load() + nb_seconds = self.time[-1] + t_span = np.linspace(0, nb_seconds, self.frames) + prr_model = prr.BiorbdModel.from_biorbd_object(self.model) + + # pyorerun animation + rerun_biorbd = prr.MultiPhaseRerun() + rerun_biorbd.add_phase(t_span=t_span, phase=0, window="animation") + rerun_biorbd.add_animated_model(prr_model, self.state_q, phase=0, window="animation") + + for path in additional_path: + with open(path, "rb") as f: + data = pickle.load(f) + + state_q = data["states"]["q"] + frames = state_q.shape[1] + + t_span = np.linspace(0, nb_seconds, frames) + prr_model = prr.BiorbdModel.from_biorbd_object(self.model) + + rerun_biorbd.add_phase(t_span=t_span, phase=0, window="split_animation") + rerun_biorbd.add_animated_model(prr_model, state_q, phase=0, window="split_animation") + + rerun_biorbd.rerun("multi_model_test") diff --git a/cocofest/result/pickle.py b/cocofest/result/pickle.py index bae45387..27a1fea1 100644 --- a/cocofest/result/pickle.py +++ b/cocofest/result/pickle.py @@ -27,6 +27,9 @@ def pickle(self): "parameters": self.sol.decision_parameters(), "parameters_bounds": bounds, "time_to_optimize": self.sol.real_time_to_optimize, + "bio_model_path": ( + self.sol.ocp.nlp[0].model.bio_model.path if hasattr(self.sol.ocp.nlp[0].model, "bio_model") else None + ), } with open(self.path + self.file_name, "wb") as file: diff --git a/environment.yml b/environment.yml index 985d78a0..f9dae504 100644 --- a/environment.yml +++ b/environment.yml @@ -7,3 +7,4 @@ dependencies: - bioviz - python-graphviz - pyqtgraph +- pyorerun