black

cocofest/examples/dynamics/reaching_task/visualization/force_and_fatigue_subplot.py

@@ -215,5 +215,5 @@
     weight="bold",
 )
 axs.text(0.75, 96.3, "Time (s)", ha="center", va="center", fontsize=18, weight="bold")
-axs.legend(title='Cost function', fontsize="medium", loc="upper right", ncol=1)
+axs.legend(title="Cost function", fontsize="medium", loc="upper right", ncol=1)
 plt.show()

cocofest/examples/getting_started/pulse_duration_optimization_nmpc_cyclic.py

@@ -14,28 +14,30 @@
 
 # --- Build target force --- #
 target_time = np.linspace(0, 1, 100)
-target_force = abs(np.sin(target_time*np.pi)) * 200
+target_force = abs(np.sin(target_time * np.pi)) * 200
 force_tracking = [target_time, target_force]
 
 # --- Build nmpc cyclic --- #
 n_total_cycles = 8
 minimum_pulse_duration = DingModelPulseDurationFrequencyWithFatigue().pd0
-nmpc = OcpFesNmpcCyclic(model=DingModelPulseDurationFrequencyWithFatigue(sum_stim_truncation=10),
-                n_stim=30,
-                n_shooting=5,
-                final_time=1,
-                pulse_duration={
-                      "min": minimum_pulse_duration,
-                      "max": 0.0006,
-                      "bimapping": False,
-                },
-                objective={"force_tracking": force_tracking},
-                n_total_cycles=n_total_cycles,
-                n_simultaneous_cycles=3,
-                n_cycle_to_advance=1,
-                cycle_to_keep="middle",
-                use_sx=True,
-                ode_solver=OdeSolver.COLLOCATION())
+nmpc = OcpFesNmpcCyclic(
+    model=DingModelPulseDurationFrequencyWithFatigue(sum_stim_truncation=10),
+    n_stim=30,
+    n_shooting=5,
+    final_time=1,
+    pulse_duration={
+        "min": minimum_pulse_duration,
+        "max": 0.0006,
+        "bimapping": False,
+    },
+    objective={"force_tracking": force_tracking},
+    n_total_cycles=n_total_cycles,
+    n_simultaneous_cycles=3,
+    n_cycle_to_advance=1,
+    cycle_to_keep="middle",
+    use_sx=True,
+    ode_solver=OdeSolver.COLLOCATION(),
+)
 
 nmpc.prepare_nmpc()
 nmpc.solve()
@@ -46,22 +48,22 @@
 force = [j for sub in nmpc.result["states"]["F"] for j in sub]
 
 ax1 = plt.subplot(221)
-ax1.plot(time, fatigue, label='A', color='green')
-ax1.set_title('Fatigue', weight='bold')
-ax1.set_xlabel('Time (s)')
-ax1.set_ylabel('Force scaling factor (-)')
+ax1.plot(time, fatigue, label="A", color="green")
+ax1.set_title("Fatigue", weight="bold")
+ax1.set_xlabel("Time (s)")
+ax1.set_ylabel("Force scaling factor (-)")
 plt.legend()
 
 ax2 = plt.subplot(222)
-ax2.plot(time, force, label='F', color='red', linewidth=4)
+ax2.plot(time, force, label="F", color="red", linewidth=4)
 for i in range(n_total_cycles):
     if i == 0:
-        ax2.plot(target_time, target_force, label='Target', color='purple')
+        ax2.plot(target_time, target_force, label="Target", color="purple")
     else:
-        ax2.plot(target_time + i, target_force, color='purple')
-ax2.set_title('Force', weight='bold')
-ax2.set_xlabel('Time (s)')
-ax2.set_ylabel('Force (N)')
+        ax2.plot(target_time + i, target_force, color="purple")
+ax2.set_title("Force", weight="bold")
+ax2.set_xlabel("Time (s)")
+ax2.set_ylabel("Force (N)")
 plt.legend()
 
 barWidth = 0.25  # set width of bar
@@ -76,11 +78,10 @@
 
 ax3 = plt.subplot(212)
 for i in range(n_total_cycles):
-    ax3.bar(bar[i], cycles[i], width = barWidth,
-            edgecolor ='grey', label =f'cycle n°{i+1}')
-ax3.set_xticks([np.mean(r) for r in bar], [str(i+1) for i in range(n_total_cycles)])
-ax3.set_xlabel('Cycles')
-ax3.set_ylabel('Pulse duration (s)')
+    ax3.bar(bar[i], cycles[i], width=barWidth, edgecolor="grey", label=f"cycle n°{i+1}")
+ax3.set_xticks([np.mean(r) for r in bar], [str(i + 1) for i in range(n_total_cycles)])
+ax3.set_xlabel("Cycles")
+ax3.set_ylabel("Pulse duration (s)")
 plt.legend()
-ax3.set_title('Pulse duration', weight='bold')
+ax3.set_title("Pulse duration", weight="bold")
 plt.show()

cocofest/models/ding2003.py

@@ -34,7 +34,7 @@ def __init__(
         model_name: str = "ding2003",
         muscle_name: str = None,
         sum_stim_truncation: int = None,
-        stim_prev: list[float] = None
+        stim_prev: list[float] = None,
     ):
         super().__init__()
         self._model_name = model_name

cocofest/models/ding2007.py

@@ -27,7 +27,13 @@ class DingModelPulseDurationFrequency(DingModelFrequency):
     Muscle & Nerve: Official Journal of the American Association of Electrodiagnostic Medicine, 36(2), 214-222.
     """
 
-    def __init__(self, model_name: str = "ding_2007", muscle_name: str = None, sum_stim_truncation: int = None, stim_prev: list[float] = None):
+    def __init__(
+        self,
+        model_name: str = "ding_2007",
+        muscle_name: str = None,
+        sum_stim_truncation: int = None,
+        stim_prev: list[float] = None,
+    ):
         super(DingModelPulseDurationFrequency, self).__init__(
             model_name=model_name, muscle_name=muscle_name, sum_stim_truncation=sum_stim_truncation, stim_prev=stim_prev
         )

cocofest/models/ding2007_with_fatigue.py

@@ -27,7 +27,11 @@ class DingModelPulseDurationFrequencyWithFatigue(DingModelPulseDurationFrequency
     """
 
     def __init__(
-        self, model_name: str = "ding_2007_with_fatigue", muscle_name: str = None, sum_stim_truncation: int = None, stim_prev: list[float] = None
+        self,
+        model_name: str = "ding_2007_with_fatigue",
+        muscle_name: str = None,
+        sum_stim_truncation: int = None,
+        stim_prev: list[float] = None,
     ):
         super(DingModelPulseDurationFrequencyWithFatigue, self).__init__(
             model_name=model_name, muscle_name=muscle_name, sum_stim_truncation=sum_stim_truncation, stim_prev=stim_prev

cocofest/optimization/fes_ocp_nmpc_cyclic.py

@@ -1,30 +1,40 @@
 import math
 
 import numpy as np
-from bioptim import SolutionMerge, ObjectiveList, ObjectiveFcn, OdeSolver, Node, OptimalControlProgram, ControlType, TimeAlignment
+from bioptim import (
+    SolutionMerge,
+    ObjectiveList,
+    ObjectiveFcn,
+    OdeSolver,
+    Node,
+    OptimalControlProgram,
+    ControlType,
+    TimeAlignment,
+)
 
 from .fes_ocp import OcpFes
 from ..models.fes_model import FesModel
 from ..custom_objectives import CustomObjective
 
 
 class OcpFesNmpcCyclic:
-    def __init__(self,
-                 model: FesModel = None,
-                 n_stim: int = None,
-                 n_shooting: int = None,
-                 final_time: int | float = None,
-                 pulse_event: dict = None,
-                 pulse_duration: dict = None,
-                 pulse_intensity: dict = None,
-                 n_total_cycles: int = None,
-                 n_simultaneous_cycles: int = None,
-                 n_cycle_to_advance: int = None,
-                 cycle_to_keep: str = None,
-                 objective: dict = None,
-                 use_sx: bool = True,
-                 ode_solver: OdeSolver = OdeSolver.RK4(n_integration_steps=1),
-                 n_threads: int = 1,
+    def __init__(
+        self,
+        model: FesModel = None,
+        n_stim: int = None,
+        n_shooting: int = None,
+        final_time: int | float = None,
+        pulse_event: dict = None,
+        pulse_duration: dict = None,
+        pulse_intensity: dict = None,
+        n_total_cycles: int = None,
+        n_simultaneous_cycles: int = None,
+        n_cycle_to_advance: int = None,
+        cycle_to_keep: str = None,
+        objective: dict = None,
+        use_sx: bool = True,
+        ode_solver: OdeSolver = OdeSolver.RK4(n_integration_steps=1),
+        n_threads: int = 1,
     ):
         super(OcpFesNmpcCyclic, self).__init__()
         self.model = model
@@ -104,7 +114,9 @@ def prepare_nmpc(self):
             n_threads=self.n_threads,
         )
 
-        OcpFes._sanity_check_frequency(n_stim=self.n_stim, final_time=self.final_time, frequency=frequency, round_down=round_down)
+        OcpFes._sanity_check_frequency(
+            n_stim=self.n_stim, final_time=self.final_time, frequency=frequency, round_down=round_down
+        )
 
         force_fourier_coefficient = (
             None if force_tracking is None else OcpFes._build_fourier_coefficient(force_tracking)
@@ -146,7 +158,14 @@ def prepare_nmpc(self):
         x_bounds, x_init = OcpFes._set_bounds(self.model, self.n_stim * self.n_simultaneous_cycles)
         one_cycle_shooting = [self.n_shooting] * self.n_stim
         objective_functions = self._set_objective(
-            self.n_stim, one_cycle_shooting, force_fourier_coefficient, end_node_tracking, custom_objective, time_min, time_max, self.n_simultaneous_cycles
+            self.n_stim,
+            one_cycle_shooting,
+            force_fourier_coefficient,
+            end_node_tracking,
+            custom_objective,
+            time_min,
+            time_max,
+            self.n_simultaneous_cycles,
         )
         all_cycle_n_shooting = [self.n_shooting] * self.n_stim * self.n_simultaneous_cycles
         self.ocp = OptimalControlProgram(
@@ -183,7 +202,7 @@ def update_stim(self, sol):
         if "pulse_apparition_time" in sol.decision_parameters():
             stimulation_time = sol.decision_parameters()["pulse_apparition_time"]
         else:
-            stimulation_time = [0] + list(np.cumsum(sol.ocp.phase_time[:self.n_stim-1]))
+            stimulation_time = [0] + list(np.cumsum(sol.ocp.phase_time[: self.n_stim - 1]))
 
         stim_prev = list(np.array(stimulation_time) - self.final_time)
         if self.previous_stim:
@@ -206,14 +225,22 @@ def store_results(self, sol_time, sol_states, sol_parameters, index):
 
             # Initialize the dict if it's the first iteration
             if index == 0:
-                self.result["time"] = [None]*self.n_total_cycles
-                [self.result["states"].update({state_key: [None]*self.n_total_cycles}) for state_key in list(sol_states[0].keys())]
-                [self.result["parameters"].update({key_parameter: [None]*self.n_total_cycles}) for key_parameter in list(sol_parameters.keys())]
+                self.result["time"] = [None] * self.n_total_cycles
+                [
+                    self.result["states"].update({state_key: [None] * self.n_total_cycles})
+                    for state_key in list(sol_states[0].keys())
+                ]
+                [
+                    self.result["parameters"].update({key_parameter: [None] * self.n_total_cycles})
+                    for key_parameter in list(sol_parameters.keys())
+                ]
 
             # Store the results
             phase_size = np.array(sol_time).shape[0]
             node_size = np.array(sol_time).shape[1]
-            sol_time = list(np.array(sol_time).reshape(phase_size*node_size))[self.first_node_in_phase*node_size:self.last_node_in_phase*node_size]
+            sol_time = list(np.array(sol_time).reshape(phase_size * node_size))[
+                self.first_node_in_phase * node_size : self.last_node_in_phase * node_size
+            ]
             sol_time = list(dict.fromkeys(sol_time))  # Remove duplicate time
             if index == 0:
                 updated_sol_time = [t - sol_time[0] for t in sol_time]
@@ -223,13 +250,17 @@ def store_results(self, sol_time, sol_states, sol_parameters, index):
             self.result["time"][index] = updated_sol_time[:-1]
 
             for state_key in list(sol_states[0].keys()):
-                middle_states_values = sol_states[self.first_node_in_phase:self.last_node_in_phase]
-                middle_states_values = [list(middle_states_values[i][state_key][0])[:-1] for i in range(len(middle_states_values))]  # Remove the last node duplicate
+                middle_states_values = sol_states[self.first_node_in_phase : self.last_node_in_phase]
+                middle_states_values = [
+                    list(middle_states_values[i][state_key][0])[:-1] for i in range(len(middle_states_values))
+                ]  # Remove the last node duplicate
                 middle_states_values = [j for sub in middle_states_values for j in sub]
                 self.result["states"][state_key][index] = middle_states_values
 
             for key_parameter in list(sol_parameters.keys()):
-                self.result["parameters"][key_parameter][index] = sol_parameters[key_parameter][self.first_node_in_phase:self.last_node_in_phase]
+                self.result["parameters"][key_parameter][index] = sol_parameters[key_parameter][
+                    self.first_node_in_phase : self.last_node_in_phase
+                ]
         return
 
     def solve(self):
@@ -244,13 +275,23 @@ def solve(self):
             # Todo uncomment when the model is updated to take into account the past stimulation
 
     @staticmethod
-    def _set_objective(n_stim, n_shooting, force_fourier_coefficient, end_node_tracking, custom_objective, time_min, time_max,
-                       n_simultaneous_cycles):
+    def _set_objective(
+        n_stim,
+        n_shooting,
+        force_fourier_coefficient,
+        end_node_tracking,
+        custom_objective,
+        time_min,
+        time_max,
+        n_simultaneous_cycles,
+    ):
         # Creates the objective for our problem
         objective_functions = ObjectiveList()
         if custom_objective:
             if len(custom_objective) != n_stim:
-                raise ValueError("The number of custom objective must be equal to the stimulation number of a single cycle")
+                raise ValueError(
+                    "The number of custom objective must be equal to the stimulation number of a single cycle"
+                )
             for i in range(len(custom_objective)):
                 for j in range(n_simultaneous_cycles):
                     objective_functions.add(custom_objective[i + j * n_stim][0])