From 86836b6d3ee8b2d404a4f49cd872dd33c9e6afc3 Mon Sep 17 00:00:00 2001 From: Pariterre Date: Tue, 3 Dec 2024 15:56:38 -0500 Subject: [PATCH 01/19] Added a casadi function interface with tests (example to come) --- bioptim/__init__.py | 4 +- bioptim/interfaces/__init__.py | 1 + .../interfaces/casadi_function_interface.py | 153 ++++++++++++++++++ .../shard5/test_casadi_function_interface.py | 110 +++++++++++++ 4 files changed, 266 insertions(+), 2 deletions(-) create mode 100644 bioptim/interfaces/casadi_function_interface.py create mode 100644 tests/shard5/test_casadi_function_interface.py diff --git a/bioptim/__init__.py b/bioptim/__init__.py index 56188df4a..8332f0e99 100644 --- a/bioptim/__init__.py +++ b/bioptim/__init__.py @@ -166,7 +166,7 @@ from .dynamics.ode_solvers import OdeSolver, OdeSolverBase from .gui.online_callback_server import PlottingServer from .gui.plot import CustomPlot -from .interfaces import Solver +from .interfaces import Solver, CasadiFunctionInterface from .limits.constraints import ConstraintFcn, ConstraintList, Constraint, ParameterConstraintList from .limits.fatigue_path_conditions import FatigueBounds, FatigueInitialGuess from .limits.multinode_constraint import MultinodeConstraintFcn, MultinodeConstraintList, MultinodeConstraint @@ -196,7 +196,7 @@ OnlineOptim, ContactType, ) -from .misc.mapping import BiMappingList, BiMapping, Mapping, SelectionMapping, Dependency +from .misc.mapping import BiMappingList, BiMapping, Mapping, NodeMapping, NodeMappingList, SelectionMapping, Dependency from .models.biorbd.biorbd_model import BiorbdModel from .models.biorbd.external_forces import ExternalForceSetTimeSeries, ExternalForceSetVariables from .models.biorbd.holonomic_biorbd_model import HolonomicBiorbdModel diff --git a/bioptim/interfaces/__init__.py b/bioptim/interfaces/__init__.py index ef615140f..221e45bae 100644 --- a/bioptim/interfaces/__init__.py +++ b/bioptim/interfaces/__init__.py @@ -2,6 +2,7 @@ from .fratrop_options import FATROP from .acados_options import ACADOS from .sqp_options import SQP_METHOD +from .casadi_function_interface import CasadiFunctionInterface class Solver: diff --git a/bioptim/interfaces/casadi_function_interface.py b/bioptim/interfaces/casadi_function_interface.py new file mode 100644 index 000000000..feb986c2f --- /dev/null +++ b/bioptim/interfaces/casadi_function_interface.py @@ -0,0 +1,153 @@ +from abc import ABC, abstractmethod + +from casadi import Callback, Function, Sparsity, DM, MX, SX +import numpy as np + + +class CasadiFunctionInterface(Callback, ABC): + def __init__(self, name: str, opts={}): + self.reverse_function = None + + super(CasadiFunctionInterface, self).__init__() + self.construct(name, opts) # Defines the self.mx_in() + self._cached_mx_in = super().mx_in() + + @abstractmethod + def inputs_len(self) -> list[int]: + """ + The len of the inputs of the function. This will help create the MX/SX vectors such that each element of the list + is the length of the input vector (i.e. the sparsity of the input vector). + + Example: + def inputs_len(self) -> list[int]: + return [3, 4] # Assuming two inputs x and y of length 3 and 4 respectively + """ + pass + + @abstractmethod + def outputs_len(self) -> list[int]: + """ + The len of the outputs of the function. This will help create the MX/SX vectors such that each element of the list + is the length of the output vector (i.e. the sparsity of the output vector). + + Example: + def outputs_len(self) -> list[int]: + return [5] # Assuming the output is a 5x1 vector + """ + pass + + @abstractmethod + def function(self, *args) -> np.ndarray | DM: + """ + The actual function to interface with casadi. The callable that returns should be callable by function(*mx_in). + If your function needs more parameters, they should be encapsulated in a partial. + + Example: + def function(self, x, y): + x = np.array(x)[:, 0] + y = np.array(y)[:, 0] + return np.array( + [ + x[0] * y[1] + x[0] * y[0] * y[0], + x[1] * x[1] + 2 * y[1], + x[0] * x[1] * x[2], + x[2] * x[1] * y[2] + 2 * y[3] * y[2], + y[0] * y[1] * y[2] * y[3], + ] + ) + """ + pass + + @abstractmethod + def jacobians(self, *args) -> list[np.ndarray | DM]: + """ + All the jacobians evaluated at *args. Each of the jacobian should be of the shape (n_out, n_in), where n_out is + the length of the output vector (the same for all) and n_in is the length of the input element (specific to each + input element). + + Example: + def jacobians(self, x, y): + x = np.array(x)[:, 0] + y = np.array(y)[:, 0] + jacobian_x = np.array( + [ + [y[1] + y[0] * y[0], 0, 0], + [0, 2 * x[1], 0], + [x[1] * x[2], x[0] * x[2], x[0] * x[1]], + [0, x[2] * y[2], x[1] * y[2]], + [0, 0, 0], + ] + ) + jacobian_y = np.array( + [ + [x[0] * 2 * y[0], x[0], 0, 0], + [0, 2, 0, 0], + [0, 0, 0, 0], + [0, 0, x[1] * x[2] + 2 * y[3], 2 * y[2]], + [y[1] * y[2] * y[3], y[0] * y[2] * y[3], y[0] * y[1] * y[3], y[0] * y[1] * y[2]], + ] + ) + return [jacobian_x, jacobian_y] # There are as many jacobians as there are inputs + """ + pass + + def mx_in(self) -> MX: + """ + Get the MX in, but it is ensured that the MX are the same at each call + """ + return self._cached_mx_in + + def get_n_in(self): + return len(self.inputs_len()) + + def get_n_out(self): + return len(self.outputs_len()) + + def get_sparsity_in(self, i): + return Sparsity.dense(self.inputs_len()[i], 1) + + def get_sparsity_out(self, i): + return Sparsity.dense(self.outputs_len()[i], 1) + + def eval(self, *args): + return [self.function(*args[0])] + + def has_reverse(self, nadj): + return nadj == 1 + + def get_reverse(self, nadj, name, inames, onames, opts): + class Reverse(Callback): + def __init__(self, parent, jacobian_functions, opts={}): + self._sparsity_in = parent.mx_in() + parent.mx_out() + self._sparsity_out = parent.mx_in() + + self.jacobian_functions = jacobian_functions + Callback.__init__(self) + self.construct("Reverse", opts) + + def get_n_in(self): + return len(self._sparsity_in) + + def get_n_out(self): + return len(self._sparsity_out) + + def get_sparsity_in(self, i): + return Sparsity.dense(self._sparsity_in[i].shape) + + def get_sparsity_out(self, i): + return Sparsity.dense(self._sparsity_out[i].shape) + + def eval(self, arg): + # Find the index to evaluate from the last parameter which is a DM vector of 0s with one value being 1 + index = arg[-1].toarray()[:, 0].tolist().index(1.0) + inputs = arg[:-1] + return [jaco[index, :].T for jaco in self.jacobian_functions(*inputs)] + + # Package it in the [nominal_in + nominal_out + adj_seed] form that CasADi expects + if self.reverse_function is None: + self.reverse_function = Reverse(self, self.jacobians) + + cx_in = self.mx_in() + nominal_out = self.mx_out() + adj_seed = self.mx_out() + return Function(name, cx_in + nominal_out + adj_seed, self.reverse_function(*cx_in, adj_seed[0])) diff --git a/tests/shard5/test_casadi_function_interface.py b/tests/shard5/test_casadi_function_interface.py new file mode 100644 index 000000000..e3e4be44e --- /dev/null +++ b/tests/shard5/test_casadi_function_interface.py @@ -0,0 +1,110 @@ +from casadi import MX, vertcat, Function, jacobian +import numpy as np +import numpy.testing as npt +from bioptim import CasadiFunctionInterface + + +class CasadiFunctionInterfaceTest(CasadiFunctionInterface): + """ + This example implements a somewhat simple 5x1 function, with x and y inputs (x => 3x1; y => 4x1) of the form + f(x, y) = np.array( + [ + x[0] * y[1] + y[0] * y[0], + x[1] * x[1] + 2 * y[1], + x[0] * x[1] * x[2], + x[2] * x[1] + 2 * y[3] * y[2], + y[0] * y[1] * y[2] * y[3], + ] + ) + + It implements the equation (5x1) and the jacobians for the inputs x (5x3) and y (5x4). + """ + + def __init__(self, opts={}): + super(CasadiFunctionInterfaceTest, self).__init__("CasadiFunctionInterfaceTest", opts) + + def inputs_len(self) -> list[int]: + return [3, 4] + + def outputs_len(self) -> list[int]: + return [5] + + def function(self, *args): + x, y = args + x = np.array(x)[:, 0] + y = np.array(y)[:, 0] + return np.array( + [ + x[0] * y[1] + x[0] * y[0] * y[0], + x[1] * x[1] + 2 * y[1], + x[0] * x[1] * x[2], + x[2] * x[1] * y[2] + 2 * y[3] * y[2], + y[0] * y[1] * y[2] * y[3], + ] + ) + + def jacobians(self, *args): + x, y = args + x = np.array(x)[:, 0] + y = np.array(y)[:, 0] + jacobian_x = np.array( + [ + [y[1] + y[0] * y[0], 0, 0], + [0, 2 * x[1], 0], + [x[1] * x[2], x[0] * x[2], x[0] * x[1]], + [0, x[2] * y[2], x[1] * y[2]], + [0, 0, 0], + ] + ) + jacobian_y = np.array( + [ + [x[0] * 2 * y[0], x[0], 0, 0], + [0, 2, 0, 0], + [0, 0, 0, 0], + [0, 0, x[1] * x[2] + 2 * y[3], 2 * y[2]], + [y[1] * y[2] * y[3], y[0] * y[2] * y[3], y[0] * y[1] * y[3], y[0] * y[1] * y[2]], + ] + ) + return [jacobian_x, jacobian_y] + + +def test_penalty_minimize_time(): + """ + These tests seem to test the interface, but actually all the internal methods are also called, which is what should + be tested. + """ + + # Computing the example + interface_test = CasadiFunctionInterfaceTest() + + # Testing the interface + npt.assert_equal(interface_test.inputs_len(), [3, 4]) + npt.assert_equal(interface_test.outputs_len(), [5]) + assert id(interface_test.mx_in()) == id(interface_test.mx_in()) # Calling twice returns the same object + + # Test the class can be called with DM + x_num = np.array([1.1, 2.3, 3.5]) + y_num = np.array([4.2, 5.4, 6.6, 7.7]) + npt.assert_almost_equal(interface_test(x_num, y_num), np.array([[25.344, 16.09, 8.855, 154.77, 1152.5976]]).T) + + # Test the jacobian is correct + x = MX.sym("x", interface_test.inputs_len()[0], 1) + y = MX.sym("y", interface_test.inputs_len()[1], 1) + jaco_x = Function("jaco_x", [x, y], [jacobian(interface_test(x, y), x)]) + jaco_y = Function("jaco_y", [x, y], [jacobian(interface_test(x, y), y)]) + + # Computing the same equations (and derivative) by casadi + real = vertcat( + x[0] * y[1] + x[0] * y[0] * y[0], + x[1] * x[1] + 2 * y[1], + x[0] * x[1] * x[2], + x[2] * x[1] * y[2] + 2 * y[3] * y[2], + y[0] * y[1] * y[2] * y[3], + ) + real_function = Function("real", [x, y], [real]) + jaco_x_real = Function("jaco_x_real", [x, y], [jacobian(real, x)]) + jaco_y_real = Function("jaco_y_real", [x, y], [jacobian(real, y)]) + + npt.assert_almost_equal(np.array(interface_test(x_num, y_num)), real_function(x_num, y_num)) + npt.assert_almost_equal(np.array(jaco_x(x_num, y_num)), jaco_x_real(x_num, y_num)) + npt.assert_almost_equal(np.array(jaco_y(x_num, y_num)), jaco_y_real(x_num, y_num)) From 22c7052a3d7684e3794b12a9fd87bceff7aac5be Mon Sep 17 00:00:00 2001 From: Pariterre Date: Thu, 5 Dec 2024 17:12:12 -0500 Subject: [PATCH 02/19] Trying to implement the stuff in bioptim --- bioptim/examples/__main__.py | 1 + .../custom_non_casadi_dynamics.py | 442 ++++++++++++++++++ .../interfaces/casadi_function_interface.py | 17 +- 3 files changed, 458 insertions(+), 2 deletions(-) create mode 100644 bioptim/examples/getting_started/custom_non_casadi_dynamics.py diff --git a/bioptim/examples/__main__.py b/bioptim/examples/__main__.py index 995472e3b..523141908 100644 --- a/bioptim/examples/__main__.py +++ b/bioptim/examples/__main__.py @@ -40,6 +40,7 @@ ("Custom Bounds", "custom_bounds.py"), ("Custom constraint", "custom_constraint.py"), ("Custom initial guess", "custom_initial_guess.py"), + ("Custom non casadi dynamics", "custom_non_casadi_dynamics.py"), ("Custom objectives", "custom_objectives.py"), ("Custom parameters", "custom_parameters.py"), ("Custom phase transitions", "custom_phase_transitions.py"), diff --git a/bioptim/examples/getting_started/custom_non_casadi_dynamics.py b/bioptim/examples/getting_started/custom_non_casadi_dynamics.py new file mode 100644 index 000000000..818b0cdbd --- /dev/null +++ b/bioptim/examples/getting_started/custom_non_casadi_dynamics.py @@ -0,0 +1,442 @@ +""" +TODO: Explain what is this example about + +This example is similar to the getting_started/pendulum.py example, but the dynamics are computed using a non-casadi +based model. This is useful when the dynamics are computed using a different library (e.g. TensorFlow, PyTorch, etc.) +""" + +import biorbd +from bioptim import ( + OptimalControlProgram, + DynamicsFcn, + Objective, + ObjectiveFcn, + BoundsList, + OdeSolver, + OdeSolverBase, + PhaseDynamics, + ControlType, + InitialGuessList, + Dynamics, + CasadiFunctionInterface, + BiorbdModel, + PenaltyController, + Node, +) +from casadi import Function, jacobian, MX, DM + +import numpy as np + + +class CasadiFunctionInterfaceTest(CasadiFunctionInterface): + """ + This example implements a somewhat simple 5x1 function, with x and y inputs (x => 3x1; y => 4x1) of the form + f(x, y) = np.array( + [ + x[0] * y[1] + y[0] * y[0], + x[1] * x[1] + 2 * y[1], + x[0] * x[1] * x[2], + x[2] * x[1] + 2 * y[3] * y[2], + y[0] * y[1] * y[2] * y[3], + ] + ) + + It implements the equation (5x1) and the jacobians for the inputs x (5x3) and y (5x4). + """ + + def __init__(self, model, opts={}): + super(CasadiFunctionInterfaceTest, self).__init__("CasadiFunctionInterfaceTest", opts) + + def inputs_len(self) -> list[int]: + return [2, 2] + + def outputs_len(self) -> list[int]: + return [2] + + def function(self, *args): + x, y = args + x = np.array(x)[:, 0] + y = np.array(y)[:, 0] + return np.array([x[0] * y[1] + x[0] * y[0] * y[0], x[1] * x[1] + 2 * y[1]]) + + def jacobians(self, *args): + x, y = args + x = np.array(x)[:, 0] + y = np.array(y)[:, 0] + jacobian_x = np.array([[y[1] + y[0] * y[0], 0, 0], [0, 2 * x[1], 0]]) + jacobian_y = np.array([[x[0] * 2 * y[0], x[0], 0, 0], [0, 2, 0, 0]]) + return [jacobian_x, jacobian_y] + + +class ForwardDynamicsInterface(CasadiFunctionInterface): + def __init__(self, model: BiorbdModel, opts={}): + self.non_casadi_model = biorbd.Model(model.path) + super(ForwardDynamicsInterface, self).__init__("ForwardDynamicsInterface", opts) + + def inputs_len(self) -> list[int]: + return [1] + + def outputs_len(self) -> list[int]: + return [1] + + def function(self, *args): + return [args[0]] # self.non_casadi_model.ForwardDynamics(*self.mx_in()[:3]) + + def jacobians(self, *args): + return [0, 0, DM(1), 0, 0] + + +def custom_func_track_markers(controller: PenaltyController) -> MX: + return controller.model.custom_interface(controller.states["q"].cx, controller.controls["tau"].cx) + + +def prepare_ocp( + biorbd_model_path: str, + final_time: float, + n_shooting: int, + ode_solver: OdeSolverBase = OdeSolver.RK4(), + use_sx: bool = False, + n_threads: int = 1, + phase_dynamics: PhaseDynamics = PhaseDynamics.SHARED_DURING_THE_PHASE, + expand_dynamics: bool = True, + control_type: ControlType = ControlType.CONSTANT, +) -> OptimalControlProgram: + """ + The initialization of an ocp + + Parameters + ---------- + biorbd_model_path: str + The path to the biorbd model + final_time: float + The time in second required to perform the task + n_shooting: int + The number of shooting points to define int the direct multiple shooting program + ode_solver: OdeSolverBase = OdeSolver.RK4() + Which type of OdeSolver to use + use_sx: bool + If the SX variable should be used instead of MX (can be extensive on RAM) + n_threads: int + The number of threads to use in the paralleling (1 = no parallel computing) + phase_dynamics: PhaseDynamics + If the dynamics equation within a phase is unique or changes at each node. + PhaseDynamics.SHARED_DURING_THE_PHASE is much faster, but lacks the capability to have changing dynamics within + a phase. A good example of when PhaseDynamics.ONE_PER_NODE should be used is when different external forces + are applied at each node + expand_dynamics: bool + If the dynamics function should be expanded. Please note, this will solve the problem faster, but will slow down + the declaration of the OCP, so it is a trade-off. Also depending on the solver, it may or may not work + (for instance IRK is not compatible with expanded dynamics) + control_type: ControlType + The type of the controls + + Returns + ------- + The OptimalControlProgram ready to be solved + """ + + bio_model = BiorbdModel(biorbd_model_path) + bio_model.custom_interface = CasadiFunctionInterfaceTest(bio_model) + + # Add objective functions + objective_functions = Objective(custom_func_track_markers, custom_type=ObjectiveFcn.Mayer, node=Node.START) + + # Dynamics + dynamics = Dynamics(DynamicsFcn.TORQUE_DRIVEN, expand_dynamics=expand_dynamics, phase_dynamics=phase_dynamics) + + # Path bounds + x_bounds = BoundsList() + x_bounds["q"] = bio_model.bounds_from_ranges("q") + x_bounds["q"][:, [0, -1]] = 0 # Start and end at 0... + x_bounds["q"][1, -1] = 3.14 # ...but end with pendulum 180 degrees rotated + x_bounds["qdot"] = bio_model.bounds_from_ranges("qdot") + x_bounds["qdot"][:, [0, -1]] = 0 # Start and end without any velocity + + # Initial guess (optional since it is 0, we show how to initialize anyway) + x_init = InitialGuessList() + x_init["q"] = [0] * bio_model.nb_q + x_init["qdot"] = [0] * bio_model.nb_qdot + + # Define control path bounds + n_tau = bio_model.nb_tau + u_bounds = BoundsList() + u_bounds["tau"] = [-100] * n_tau, [100] * n_tau # Limit the strength of the pendulum to (-100 to 100)... + u_bounds["tau"][1, :] = 0 # ...but remove the capability to actively rotate + + # Initial guess (optional since it is 0, we show how to initialize anyway) + u_init = InitialGuessList() + u_init["tau"] = [0] * n_tau + + return OptimalControlProgram( + bio_model, + dynamics, + n_shooting, + final_time, + x_init=x_init, + u_init=u_init, + x_bounds=x_bounds, + u_bounds=u_bounds, + objective_functions=objective_functions, + ode_solver=ode_solver, + use_sx=use_sx, + n_threads=n_threads, + control_type=control_type, + ) + + +def main(): + """ + If pendulum is run as a script, it will perform the optimization and animates it + """ + + # --- Prepare the ocp --- # + ocp = prepare_ocp(biorbd_model_path="models/pendulum.bioMod", final_time=1, n_shooting=400, n_threads=2) + + # --- Solve the ocp --- # + sol = ocp.solve() + + # --- Show the results graph --- # + # sol.print_cost() + sol.graphs(show_bounds=True, save_name="results.png") + + +if __name__ == "__main__": + main() + + +######## OCP FAST ######## +# from casadi import * + +# T = 10.0 # Time horizon +# N = 20 # number of control intervals + +# # Declare model variables +# x1 = MX.sym("x1") +# x2 = MX.sym("x2") +# x = vertcat(x1, x2) +# u = MX.sym("u") + +# # Model equations +# xdot = vertcat((1 - x2**2) * x1 - x2 + u, x1) + + +# # Formulate discrete time dynamics +# if False: +# # CVODES from the SUNDIALS suite +# dae = {"x": x, "p": u, "ode": xdot} +# opts = {"tf": T / N} +# F = integrator("F", "cvodes", dae, opts) +# else: +# # Fixed step Runge-Kutta 4 integrator +# M = 4 # RK4 steps per interval +# DT = T / N / M +# f = Function("f", [x, u], [xdot]) +# X0 = MX.sym("X0", 2) +# U = MX.sym("U") +# X = X0 +# Q = 0 +# for j in range(M): +# k1 = f(X, U) +# k2 = f(X + DT / 2 * k1, U) +# k3 = f(X + DT / 2 * k2, U) +# k4 = f(X + DT * k3, U) +# X = X + DT / 6 * (k1 + 2 * k2 + 2 * k3 + k4) +# F = Function("F", [X0, U], [X], ["x0", "p"], ["xf"]) + +# # Start with an empty NLP +# w = [] +# w0 = [] +# lbw = [] +# ubw = [] +# g = [] +# lbg = [] +# ubg = [] + +# # "Lift" initial conditions +# Xk = MX.sym("X0", 2) +# w += [Xk] +# lbw += [0, 1] +# ubw += [0, 1] +# w0 += [0, 1] + +# # Formulate the NLP +# for k in range(N): +# # New NLP variable for the control +# Uk = MX.sym("U_" + str(k)) +# w += [Uk] +# lbw += [-1] +# ubw += [1] +# w0 += [0] + +# # Integrate till the end of the interval +# Fk = F(x0=Xk, p=Uk) +# Xk_end = Fk["xf"] + +# # New NLP variable for state at end of interval +# Xk = MX.sym("X_" + str(k + 1), 2) +# w += [Xk] +# lbw += [-0.25, -inf] +# ubw += [inf, inf] +# w0 += [0, 0] + +# # Add equality constraint +# g += [Xk_end - Xk] +# lbg += [0, 0] +# ubg += [0, 0] + +# nd = N + 1 + +# import gpflow +# import time + +# from tensorflow_casadi import TensorFlowEvaluator + + +# class GPR(TensorFlowEvaluator): +# def __init__(self, session, opts={}): +# X = tf.compat.v1.placeholder(shape=(1, nd), dtype=np.float64) +# mean = tf.reshape(tf.reduce_mean(X), (1, 1)) +# TensorFlowEvaluator.__init__(self, [X], [mean], session, opts) +# self.counter = 0 +# self.time = 0 + +# def eval(self, arg): +# self.counter += 1 +# t0 = time.time() +# ret = TensorFlowEvaluator.eval(self, arg) +# self.time += time.time() - t0 +# return [ret] + + +# import tensorflow as tf + +# with tf.compat.v1.Session() as session: +# GPR = GPR(session) + +# w = vertcat(*w) + +# # Create an NLP solver +# prob = {"f": sum1(GPR(w[0::3])), "x": w, "g": vertcat(*g)} +# options = {"ipopt": {"hessian_approximation": "limited-memory"}} +# solver = nlpsol("solver", "ipopt", prob, options) + +# # Solve the NLP +# sol = solver(x0=w0, lbx=lbw, ubx=ubw, lbg=lbg, ubg=ubg) + +# print("Ncalls", GPR.counter) +# print("Total time [s]", GPR.time) +# w_opt = sol["x"].full().flatten() + +# # Plot the solution +# x1_opt = w_opt[0::3] +# x2_opt = w_opt[1::3] +# u_opt = w_opt[2::3] + +# tgrid = [T / N * k for k in range(N + 1)] +# import matplotlib.pyplot as plt + +# plt.figure(1) +# plt.clf() +# plt.plot(tgrid, x1_opt, "--") +# plt.plot(tgrid, x2_opt, "-") +# plt.step(tgrid, vertcat(DM.nan(1), u_opt), "-.") +# plt.xlabel("t") +# plt.legend(["x1", "x2", "u"]) +# plt.grid() +# plt.show() + + +# +# +# +######### TENSORFLOW CASADI ######### +# import casadi +# import tensorflow as tf + + +# class TensorFlowEvaluator(casadi.Callback): +# def __init__(self, t_in, t_out, session, opts={}): +# """ +# t_in: list of inputs (tensorflow placeholders) +# t_out: list of outputs (tensors dependeant on those placeholders) +# session: a tensorflow session +# """ +# casadi.Callback.__init__(self) +# assert isinstance(t_in, list) +# self.t_in = t_in +# assert isinstance(t_out, list) +# self.t_out = t_out +# self.construct("TensorFlowEvaluator", opts) +# self.session = session +# self.refs = [] + +# def get_n_in(self): +# return len(self.t_in) + +# def get_n_out(self): +# return len(self.t_out) + +# def get_sparsity_in(self, i): +# return casadi.Sparsity.dense(*self.t_in[i].get_shape().as_list()) + +# def get_sparsity_out(self, i): +# return casadi.Sparsity.dense(*self.t_out[i].get_shape().as_list()) + +# def eval(self, arg): +# # Associate each tensorflow input with the numerical argument passed by CasADi +# d = dict((v, arg[i].toarray()) for i, v in enumerate(self.t_in)) +# # Evaluate the tensorflow expressions +# ret = self.session.run(self.t_out, feed_dict=d) +# return ret + +# # Vanilla tensorflow offers just the reverse mode AD +# def has_reverse(self, nadj): +# return nadj == 1 + +# def get_reverse(self, nadj, name, inames, onames, opts): +# # Construct tensorflow placeholders for the reverse seeds +# adj_seed = [ +# tf.compat.v1.placeholder(shape=self.sparsity_out(i).shape, dtype=tf.float64) for i in range(self.n_out()) +# ] +# # Construct the reverse tensorflow graph through 'gradients' +# grad = tf.gradients(self.t_out, self.t_in, grad_ys=adj_seed) +# # Create another TensorFlowEvaluator object +# callback = TensorFlowEvaluator(self.t_in + adj_seed, grad, self.session) +# # Make sure you keep a reference to it +# self.refs.append(callback) + +# # Package it in the nominal_in+nominal_out+adj_seed form that CasADi expects +# nominal_in = self.mx_in() +# nominal_out = self.mx_out() +# adj_seed = self.mx_out() +# return casadi.Function( +# name, nominal_in + nominal_out + adj_seed, callback.call(nominal_in + adj_seed), inames, onames +# ) + + +# if __name__ == "__main__": +# from casadi import * + +# a = tf.compat.v1.placeholder(shape=(2, 2), dtype=tf.float64) +# b = tf.compat.v1.placeholder(shape=(2, 1), dtype=tf.float64) + +# y = tf.matmul(tf.sin(a), b) + +# with tf.compat.v1.Session() as session: +# f_tf = TensorFlowEvaluator([a, b], [y], session) + +# a = MX.sym("a", 2, 2) +# b = MX.sym("a", 2, 1) +# y = f_tf(a, b) +# yref = mtimes(sin(a), b) + +# f = Function("f", [a, b], [y]) +# fref = Function("f", [a, b], [yref]) + +# print(f(DM([[1, 2], [3, 4]]), DM([[1], [3]]))) +# print(fref(DM([[1, 2], [3, 4]]), DM([[1], [3]]))) + +# f = Function("f", [a, b], [jacobian(y, a)]) +# fref = Function("f", [a, b], [jacobian(yref, a)]) +# print(f(DM([[1, 2], [3, 4]]), DM([[1], [3]]))) +# print(fref(DM([[1, 2], [3, 4]]), DM([[1], [3]]))) diff --git a/bioptim/interfaces/casadi_function_interface.py b/bioptim/interfaces/casadi_function_interface.py index feb986c2f..65d1954cc 100644 --- a/bioptim/interfaces/casadi_function_interface.py +++ b/bioptim/interfaces/casadi_function_interface.py @@ -1,6 +1,6 @@ from abc import ABC, abstractmethod -from casadi import Callback, Function, Sparsity, DM, MX, SX +from casadi import Callback, Function, Sparsity, DM, MX, jacobian import numpy as np @@ -122,6 +122,7 @@ def __init__(self, parent, jacobian_functions, opts={}): self._sparsity_out = parent.mx_in() self.jacobian_functions = jacobian_functions + self.reverse_function = None Callback.__init__(self) self.construct("Reverse", opts) @@ -143,6 +144,18 @@ def eval(self, arg): inputs = arg[:-1] return [jaco[index, :].T for jaco in self.jacobian_functions(*inputs)] + def has_reverse(self, nadj): + return nadj == 1 + + def get_reverse(self, nadj, name, inames, onames, opts): + if self.reverse_function is None: + self.reverse_function = Reverse(self, jacobian(self.jacobian_functions)) + + cx_in = self.mx_in() + nominal_out = self.mx_out() + adj_seed = self.mx_out() + return Function(name, cx_in + nominal_out + adj_seed, self.reverse_function.call(cx_in + adj_seed)) + # Package it in the [nominal_in + nominal_out + adj_seed] form that CasADi expects if self.reverse_function is None: self.reverse_function = Reverse(self, self.jacobians) @@ -150,4 +163,4 @@ def eval(self, arg): cx_in = self.mx_in() nominal_out = self.mx_out() adj_seed = self.mx_out() - return Function(name, cx_in + nominal_out + adj_seed, self.reverse_function(*cx_in, adj_seed[0])) + return Function(name, cx_in + nominal_out + adj_seed, self.reverse_function.call(cx_in + adj_seed)) From 86d45b3d94108841074e6bb0e61703f7a981f8be Mon Sep 17 00:00:00 2001 From: Pariterre Date: Mon, 9 Dec 2024 16:55:53 -0500 Subject: [PATCH 03/19] moving to torch instead --- bioptim/models/torch/torch_model.py | 1071 +++++++++++++++++++++++++++ 1 file changed, 1071 insertions(+) create mode 100644 bioptim/models/torch/torch_model.py diff --git a/bioptim/models/torch/torch_model.py b/bioptim/models/torch/torch_model.py new file mode 100644 index 000000000..0180bd35a --- /dev/null +++ b/bioptim/models/torch/torch_model.py @@ -0,0 +1,1071 @@ +from typing import Callable + +from casadi import SX, MX, vertcat, horzcat, norm_fro, Function +import numpy as np +import torch + +""" +INSTALLATION: +First, make sure pytorch is installed + + pip install torch>=2.0 --index-url https://download.pytorch.org/whl/cpu/torch_stable.html + +Then, install l4casadi as the interface between CasADi and PyTorch + + pip install l4casadi --no-build-isolation + +""" + + +class TorchModel: + """ + This class wraps a pytorch model and allows the user to call some useful functions on it. + """ + + def __init__(self, model: str | torch.nn.Module): + if not isinstance(bio_model, str) and not isinstance(bio_model, biorbd.Model): + raise ValueError("The model should be of type 'str' or 'biorbd.Model'") + + self.model = biorbd.Model(bio_model) if isinstance(bio_model, str) else bio_model + if parameters is not None: + for param_key in parameters: + parameters[param_key].apply_parameter(self) + self._friction_coefficients = friction_coefficients + + self.external_force_set = ( + self._set_external_force_set(external_force_set) if external_force_set is not None else None + ) + self._symbolic_variables() + self.biorbd_external_forces_set = self._dispatch_forces() if external_force_set else None + + # TODO: remove mx (the MX parameters should be created inside the BiorbdModel) + self.parameters = parameters.mx if parameters else MX() + + def _symbolic_variables(self): + """Declaration of MX variables of the right shape for the creation of CasADi Functions""" + self.q = MX.sym("q_mx", self.nb_q, 1) + self.qdot = MX.sym("qdot_mx", self.nb_qdot, 1) + self.qddot = MX.sym("qddot_mx", self.nb_qddot, 1) + self.qddot_joints = MX.sym("qddot_joints_mx", self.nb_qddot - self.nb_root, 1) + self.tau = MX.sym("tau_mx", self.nb_tau, 1) + self.muscle = MX.sym("muscle_mx", self.nb_muscles, 1) + self.activations = MX.sym("activations_mx", self.nb_muscles, 1) + self.external_forces = MX.sym( + "external_forces_mx", + self.external_force_set.nb_external_forces_components if self.external_force_set else 0, + 1, + ) + + def _set_external_force_set(self, external_force_set: ExternalForceSetTimeSeries): + """ + It checks the external forces and binds them to the model. + """ + external_force_set._check_segment_names(tuple([s.name().to_string() for s in self.model.segments()])) + external_force_set._check_all_string_points_of_application(self.marker_names) + external_force_set._bind() + + return external_force_set + + @property + def name(self) -> str: + # parse the path and split to get the .bioMod name + return self.model.path().absolutePath().to_string().split("/")[-1] + + @property + def path(self) -> str: + return self.model.path().relativePath().to_string() + + def copy(self): + return BiorbdModel(self.path) + + def serialize(self) -> tuple[Callable, dict]: + return BiorbdModel, dict(bio_model=self.path) + + @property + def friction_coefficients(self) -> MX | SX | np.ndarray: + return self._friction_coefficients + + def set_friction_coefficients(self, new_friction_coefficients) -> None: + if np.any(new_friction_coefficients < 0): + raise ValueError("Friction coefficients must be positive") + return self._friction_coefficients + + @property + def gravity(self) -> Function: + """ + Returns the gravity of the model. + Since the gravity is self-defined in the model, you need to provide the type of the output when calling the function like this: + model.gravity()(MX() / SX()) + """ + biorbd_return = self.model.getGravity().to_mx() + casadi_fun = Function( + "gravity", + [self.parameters], + [biorbd_return], + ["parameters"], + ["gravity"], + ) + return casadi_fun + + def set_gravity(self, new_gravity) -> None: + self.model.setGravity(new_gravity) + return + + @property + def nb_tau(self) -> int: + return self.model.nbGeneralizedTorque() + + @property + def nb_segments(self) -> int: + return self.model.nbSegment() + + def segment_index(self, name) -> int: + return biorbd.segment_index(self.model, name) + + @property + def nb_quaternions(self) -> int: + return self.model.nbQuat() + + @property + def nb_dof(self) -> int: + return self.model.nbDof() + + @property + def nb_q(self) -> int: + return self.model.nbQ() + + @property + def nb_qdot(self) -> int: + return self.model.nbQdot() + + @property + def nb_qddot(self) -> int: + return self.model.nbQddot() + + @property + def nb_root(self) -> int: + return self.model.nbRoot() + + @property + def segments(self) -> tuple[biorbd.Segment]: + return self.model.segments() + + def rotation_matrix_to_euler_angles(self, sequence: str) -> Function: + """ + Returns the rotation matrix to euler angles function. + """ + r = MX.sym("r_mx", 3, 3) + r_matrix = biorbd.Rotation(r[0, 0], r[0, 1], r[0, 2], r[1, 0], r[1, 1], r[1, 2], r[2, 0], r[2, 1], r[2, 2]) + biorbd_return = biorbd.Rotation.toEulerAngles(r_matrix, sequence).to_mx() + casadi_fun = Function( + "rotation_matrix_to_euler_angles", + [r], + [biorbd_return], + ["Rotation matrix"], + ["Euler angles"], + ) + return casadi_fun + + def homogeneous_matrices_in_global(self, segment_index: int, inverse=False) -> Function: + """ + Returns the roto-translation matrix of the segment in the global reference frame. + """ + q_biorbd = GeneralizedCoordinates(self.q) + jcs = self.model.globalJCS(q_biorbd, segment_index) + biorbd_return = jcs.transpose().to_mx() if inverse else jcs.to_mx() + casadi_fun = Function( + "homogeneous_matrices_in_global", + [self.q, self.parameters], + [biorbd_return], + ["q", "parameters"], + ["Joint coordinate system RT matrix in global"], + ) + return casadi_fun + + def homogeneous_matrices_in_child(self, segment_id) -> Function: + """ + Returns the roto-translation matrix of the segment in the child reference frame. + Since the homogeneous matrix is self-defined in the model, you need to provide the type of the output when calling the function like this: + model.homogeneous_matrices_in_child(segment_id)(MX() / SX()) + """ + biorbd_return = self.model.localJCS(segment_id).to_mx() + casadi_fun = Function( + "homogeneous_matrices_in_child", + [self.parameters], + [biorbd_return], + ["parameters"], + ["Joint coordinate system RT matrix in local"], + ) + return casadi_fun + + @property + def mass(self) -> Function: + """ + Returns the mass of the model. + Since the mass is self-defined in the model, you need to provide the type of the output when calling the function like this: + model.mass()(MX() / SX()) + """ + biorbd_return = self.model.mass().to_mx() + casadi_fun = Function( + "mass", + [self.parameters], + [biorbd_return], + ["parameters"], + ["mass"], + ) + return casadi_fun + + def rt(self, rt_index) -> Function: + q_biorbd = GeneralizedCoordinates(self.q) + biorbd_return = self.model.RT(q_biorbd, rt_index).to_mx() + casadi_fun = Function( + "rt", + [self.q, self.parameters], + [biorbd_return], + ["q", "parameters"], + ["RT matrix"], + ) + return casadi_fun + + def center_of_mass(self) -> Function: + q_biorbd = GeneralizedCoordinates(self.q) + biorbd_return = self.model.CoM(q_biorbd, True).to_mx() + casadi_fun = Function( + "center_of_mass", + [self.q, self.parameters], + [biorbd_return], + ["q", "parameters"], + ["Center of mass"], + ) + return casadi_fun + + def center_of_mass_velocity(self) -> Function: + q_biorbd = GeneralizedCoordinates(self.q) + qdot_biorbd = GeneralizedVelocity(self.qdot) + biorbd_return = self.model.CoMdot(q_biorbd, qdot_biorbd, True).to_mx() + casadi_fun = Function( + "center_of_mass_velocity", + [self.q, self.qdot, self.parameters], + [biorbd_return], + ["q", "qdot", "parameters"], + ["Center of mass velocity"], + ) + return casadi_fun + + def center_of_mass_acceleration(self) -> Function: + q_biorbd = GeneralizedCoordinates(self.q) + qdot_biorbd = GeneralizedVelocity(self.qdot) + qddot_biorbd = GeneralizedAcceleration(self.qddot) + biorbd_return = self.model.CoMddot(q_biorbd, qdot_biorbd, qddot_biorbd, True).to_mx() + casadi_fun = Function( + "center_of_mass_acceleration", + [self.q, self.qdot, self.qddot, self.parameters], + [biorbd_return], + ["q", "qdot", "qddot", "parameters"], + ["Center of mass acceleration"], + ) + return casadi_fun + + def body_rotation_rate(self) -> Function: + q_biorbd = GeneralizedCoordinates(self.q) + qdot_biorbd = GeneralizedVelocity(self.qdot) + biorbd_return = self.model.bodyAngularVelocity(q_biorbd, qdot_biorbd, True).to_mx() + casadi_fun = Function( + "body_rotation_rate", + [self.q, self.qdot, self.parameters], + [biorbd_return], + ["q", "qdot", "parameters"], + ["Body rotation rate"], + ) + return casadi_fun + + def mass_matrix(self) -> Function: + q_biorbd = GeneralizedCoordinates(self.q) + biorbd_return = self.model.massMatrix(q_biorbd).to_mx() + casadi_fun = Function( + "mass_matrix", + [self.q, self.parameters], + [biorbd_return], + ["q", "parameters"], + ["Mass matrix"], + ) + return casadi_fun + + def non_linear_effects(self) -> Function: + q_biorbd = GeneralizedCoordinates(self.q) + qdot_biorbd = GeneralizedVelocity(self.qdot) + biorbd_return = self.model.NonLinearEffect(q_biorbd, qdot_biorbd).to_mx() + casadi_fun = Function( + "non_linear_effects", + [self.q, self.qdot, self.parameters], + [biorbd_return], + ["q", "qdot", "parameters"], + ["Non linear effects"], + ) + return casadi_fun + + def angular_momentum(self) -> Function: + q_biorbd = GeneralizedCoordinates(self.q) + qdot_biorbd = GeneralizedVelocity(self.qdot) + biorbd_return = self.model.angularMomentum(q_biorbd, qdot_biorbd, True).to_mx() + casadi_fun = Function( + "angular_momentum", + [self.q, self.qdot, self.parameters], + [biorbd_return], + ["q", "qdot", "parameters"], + ["Angular momentum"], + ) + return casadi_fun + + def reshape_qdot(self, k_stab=1) -> Function: + biorbd_return = self.model.computeQdot( + GeneralizedCoordinates(self.q), + GeneralizedCoordinates(self.qdot), # mistake in biorbd + k_stab, + ).to_mx() + casadi_fun = Function( + "reshape_qdot", + [self.q, self.qdot, self.parameters], + [biorbd_return], + ["q", "qdot", "parameters"], + ["Reshaped qdot"], + ) + return casadi_fun + + def segment_angular_velocity(self, idx) -> Function: + """ + Returns the angular velocity of the segment in the global reference frame. + """ + q_biorbd = GeneralizedCoordinates(self.q) + qdot_biorbd = GeneralizedVelocity(self.qdot) + biorbd_return = self.model.segmentAngularVelocity(q_biorbd, qdot_biorbd, idx, True).to_mx() + casadi_fun = Function( + "segment_angular_velocity", + [self.q, self.qdot, self.parameters], + [biorbd_return], + ["q", "qdot", "parameters"], + ["Segment angular velocity"], + ) + return casadi_fun + + def segment_orientation(self, idx: int, sequence: str = "xyz") -> Function: + """ + Returns the angular position of the segment in the global reference frame. + """ + q_biorbd = GeneralizedCoordinates(self.q) + rotation_matrix = self.homogeneous_matrices_in_global(idx)(q_biorbd, self.parameters)[:3, :3] + biorbd_return = self.rotation_matrix_to_euler_angles(sequence=sequence)(rotation_matrix) + casadi_fun = Function( + "segment_orientation", + [self.q, self.parameters], + [biorbd_return], + ["q", "parameters"], + ["Segment orientation"], + ) + return casadi_fun + + @property + def name_dof(self) -> tuple[str, ...]: + return tuple(s.to_string() for s in self.model.nameDof()) + + @property + def contact_names(self) -> tuple[str, ...]: + return tuple(s.to_string() for s in self.model.contactNames()) + + @property + def nb_soft_contacts(self) -> int: + return self.model.nbSoftContacts() + + @property + def soft_contact_names(self) -> tuple[str, ...]: + return self.model.softContactNames() + + def soft_contact(self, soft_contact_index, *args): + return self.model.softContact(soft_contact_index, *args) + + @property + def muscle_names(self) -> tuple[str, ...]: + return tuple(s.to_string() for s in self.model.muscleNames()) + + @property + def nb_muscles(self) -> int: + return self.model.nbMuscles() + + def torque(self) -> Function: + """ + Returns the torque from the torque_activations. + Note that tau_activation should be between 0 and 1. + """ + q_biorbd = GeneralizedCoordinates(self.q) + qdot_biorbd = GeneralizedVelocity(self.qdot) + tau_activations_biorbd = self.tau + biorbd_return = self.model.torque(tau_activations_biorbd, q_biorbd, qdot_biorbd).to_mx() + casadi_fun = Function( + "torque_activation", + [self.tau, self.q, self.qdot, self.parameters], + [biorbd_return], + ["tau", "q", "qdot", "parameters"], + ["Torque from tau activations"], + ) + return casadi_fun + + def forward_dynamics_free_floating_base(self) -> Function: + q_biorbd = GeneralizedCoordinates(self.q) + qdot_biorbd = GeneralizedVelocity(self.qdot) + qddot_joints_biorbd = GeneralizedAcceleration(self.qddot_joints) + biorbd_return = self.model.ForwardDynamicsFreeFloatingBase(q_biorbd, qdot_biorbd, qddot_joints_biorbd).to_mx() + casadi_fun = Function( + "forward_dynamics_free_floating_base", + [self.q, self.qdot, self.qddot_joints, self.parameters], + [biorbd_return], + ["q", "qdot", "qddot_joints", "parameters"], + ["qddot_root and qddot_joints"], + ) + return casadi_fun + + @staticmethod + def reorder_qddot_root_joints(qddot_root, qddot_joints) -> MX | SX: + return vertcat(qddot_root, qddot_joints) + + def _dispatch_forces(self) -> biorbd.ExternalForceSet: + """Dispatch the symbolic MX into the biorbd external forces object""" + biorbd_external_forces = self.model.externalForceSet() + + # "type of external force": (function to call, number of force components) + force_mapping = { + "in_global": (_add_global_force, 6), + "torque_in_global": (_add_torque_global, 3), + "translational_in_global": (_add_translational_global, 3), + "in_local": (_add_local_force, 6), + "torque_in_local": (_add_torque_local, 3), + } + + symbolic_counter = 0 + for force_type, val in force_mapping.items(): + add_force_func, num_force_components = val + symbolic_counter = self._dispatch_forces_of_type( + force_type, add_force_func, num_force_components, symbolic_counter, biorbd_external_forces + ) + + return biorbd_external_forces + + def _dispatch_forces_of_type( + self, + force_type: str, + add_force_func: "Callable", + num_force_components: int, + symbolic_counter: int, + biorbd_external_forces: "biorbd.ExternalForces", + ) -> int: + """ + Helper method to dispatch forces of a specific external forces. + + Parameters + ---------- + force_type: str + The type of external force to dispatch among in_global, torque_in_global, translational_in_global, in_local, torque_in_local. + add_force_func: Callable + The function to call to add the force to the biorbd external forces object. + num_force_components: int + The number of force components for the given type + symbolic_counter: int + The current symbolic counter to slice the whole external_forces mx. + biorbd_external_forces: biorbd.ExternalForces + The biorbd external forces object to add the forces to. + + Returns + ------- + int + The updated symbolic counter. + """ + for segment, forces_on_segment in getattr(self.external_force_set, force_type).items(): + for force in forces_on_segment: + force_slicer = slice(symbolic_counter, symbolic_counter + num_force_components) + + point_of_application_mx = self._get_point_of_application(force, force_slicer.stop) + + add_force_func( + biorbd_external_forces, segment, self.external_forces[force_slicer], point_of_application_mx + ) + symbolic_counter = force_slicer.stop + ( + 3 if isinstance(force["point_of_application"], np.ndarray) else 0 + ) + + return symbolic_counter + + def _get_point_of_application(self, force, stop_index) -> biorbd.NodeSegment | np.ndarray | None: + """ + Determine the point of application mx slice based on its type. Only sliced if an array is stored + + Parameters + ---------- + force : dict + The force dictionary with details on the point of application. + stop_index : int + Index position in MX where the point of application components start. + + Returns + ------- + biorbd.NodeSegment | np.ndarray | None + Returns a slice of MX, a marker node, or None if no point of application is defined. + """ + if isinstance(force["point_of_application"], np.ndarray): + return self.external_forces[slice(stop_index, stop_index + 3)] + elif isinstance(force["point_of_application"], str): + return self.model.marker(self.marker_index(force["point_of_application"])) + return None + + def forward_dynamics(self, with_contact: bool = False) -> Function: + + q_biorbd = GeneralizedCoordinates(self.q) + qdot_biorbd = GeneralizedVelocity(self.qdot) + tau_biorbd = GeneralizedTorque(self.tau) + + if with_contact: + if self.external_force_set is None: + biorbd_return = self.model.ForwardDynamicsConstraintsDirect(q_biorbd, qdot_biorbd, tau_biorbd).to_mx() + else: + biorbd_return = self.model.ForwardDynamicsConstraintsDirect( + q_biorbd, qdot_biorbd, tau_biorbd, self.biorbd_external_forces_set + ).to_mx() + casadi_fun = Function( + "constrained_forward_dynamics", + [self.q, self.qdot, self.tau, self.external_forces, self.parameters], + [biorbd_return], + ["q", "qdot", "tau", "external_forces", "parameters"], + ["qddot"], + ) + else: + if self.external_force_set is None: + biorbd_return = self.model.ForwardDynamics(q_biorbd, qdot_biorbd, tau_biorbd).to_mx() + else: + biorbd_return = self.model.ForwardDynamics( + q_biorbd, qdot_biorbd, tau_biorbd, self.biorbd_external_forces_set + ).to_mx() + casadi_fun = Function( + "forward_dynamics", + [self.q, self.qdot, self.tau, self.external_forces, self.parameters], + [biorbd_return], + ["q", "qdot", "tau", "external_forces", "parameters"], + ["qddot"], + ) + return casadi_fun + + def inverse_dynamics(self, with_contact: bool = False) -> Function: + + if with_contact: + raise NotImplementedError("Inverse dynamics with contact is not implemented yet") + + q_biorbd = GeneralizedCoordinates(self.q) + qdot_biorbd = GeneralizedVelocity(self.qdot) + qddot_biorbd = GeneralizedAcceleration(self.qddot) + if self.external_force_set is None: + biorbd_return = self.model.InverseDynamics(q_biorbd, qdot_biorbd, qddot_biorbd).to_mx() + else: + biorbd_return = self.model.InverseDynamics( + q_biorbd, qdot_biorbd, qddot_biorbd, self.biorbd_external_forces_set + ).to_mx() + casadi_fun = Function( + "inverse_dynamics", + [self.q, self.qdot, self.qddot, self.external_forces, self.parameters], + [biorbd_return], + ["q", "qdot", "qddot", "external_forces", "parameters"], + ["tau"], + ) + return casadi_fun + + def contact_forces_from_constrained_forward_dynamics(self) -> Function: + + q_biorbd = GeneralizedCoordinates(self.q) + qdot_biorbd = GeneralizedVelocity(self.qdot) + tau_biorbd = GeneralizedTorque(self.tau) + if self.external_force_set is None: + biorbd_return = self.model.ContactForcesFromForwardDynamicsConstraintsDirect( + q_biorbd, qdot_biorbd, tau_biorbd + ).to_mx() + else: + biorbd_return = self.model.ContactForcesFromForwardDynamicsConstraintsDirect( + q_biorbd, qdot_biorbd, tau_biorbd, self.biorbd_external_forces_set + ).to_mx() + casadi_fun = Function( + "contact_forces_from_constrained_forward_dynamics", + [self.q, self.qdot, self.tau, self.external_forces, self.parameters], + [biorbd_return], + ["q", "qdot", "tau", "external_forces", "parameters"], + ["contact_forces"], + ) + return casadi_fun + + def qdot_from_impact(self) -> Function: + q_biorbd = GeneralizedCoordinates(self.q) + qdot_pre_impact_biorbd = GeneralizedVelocity(self.qdot) + biorbd_return = self.model.ComputeConstraintImpulsesDirect(q_biorbd, qdot_pre_impact_biorbd).to_mx() + casadi_fun = Function( + "qdot_from_impact", + [self.q, self.qdot, self.parameters], + [biorbd_return], + ["q", "qdot", "parameters"], + ["qdot post impact"], + ) + return casadi_fun + + def muscle_activation_dot(self) -> Function: + muscle_states = self.model.stateSet() + for k in range(self.model.nbMuscles()): + muscle_states[k].setActivation(self.activations[k]) + muscle_states[k].setExcitation(self.muscle[k]) + biorbd_return = self.model.activationDot(muscle_states).to_mx() + casadi_fun = Function( + "muscle_activation_dot", + [self.muscle, self.activations, self.parameters], + [biorbd_return], + ["muscle_excitation", "muscle_activation", "parameters"], + ["muscle_activation_dot"], + ) + return casadi_fun + + def muscle_length_jacobian(self) -> Function: + q_biorbd = GeneralizedCoordinates(self.q) + biorbd_return = self.model.musclesLengthJacobian(q_biorbd).to_mx() + casadi_fun = Function( + "muscle_length_jacobian", + [self.q, self.parameters], + [biorbd_return], + ["q", "parameters"], + ["muscle_length_jacobian"], + ) + return casadi_fun + + def muscle_velocity(self) -> Function: + J = self.muscle_length_jacobian()(self.q, self.parameters) + biorbd_return = J @ self.qdot + casadi_fun = Function( + "muscle_velocity", + [self.q, self.qdot, self.parameters], + [biorbd_return], + ["q", "qdot", "parameters"], + ["muscle_velocity"], + ) + return casadi_fun + + def muscle_joint_torque(self) -> Function: + muscles_states = self.model.stateSet() + muscles_activations = self.muscle + for k in range(self.model.nbMuscles()): + muscles_states[k].setActivation(muscles_activations[k]) + q_biorbd = GeneralizedCoordinates(self.q) + qdot_biorbd = GeneralizedVelocity(self.qdot) + biorbd_return = self.model.muscularJointTorque(muscles_states, q_biorbd, qdot_biorbd).to_mx() + casadi_fun = Function( + "muscle_joint_torque", + [self.muscle, self.q, self.qdot, self.parameters], + [biorbd_return], + ["muscle_activation", "q", "qdot", "parameters"], + ["muscle_joint_torque"], + ) + return casadi_fun + + def markers(self) -> list[MX]: + biorbd_return = horzcat(*[m.to_mx() for m in self.model.markers(GeneralizedCoordinates(self.q))]) + casadi_fun = Function( + "markers", + [self.q, self.parameters], + [biorbd_return], + ["q", "parameters"], + ["markers"], + ) + return casadi_fun + + @property + def nb_markers(self) -> int: + return self.model.nbMarkers() + + def marker_index(self, name): + return biorbd.marker_index(self.model, name) + + def marker(self, index: int, reference_segment_index: int = None) -> Function: + q_biorbd = GeneralizedCoordinates(self.q) + marker = self.model.marker(q_biorbd, index) + if reference_segment_index is not None: + global_homogeneous_matrix = self.model.globalJCS(q_biorbd, reference_segment_index) + marker_rotated = global_homogeneous_matrix.transpose().to_mx() @ vertcat(marker.to_mx(), 1) + biorbd_return = marker_rotated[:3] + else: + biorbd_return = marker.to_mx() + casadi_fun = Function( + "marker", + [self.q, self.parameters], + [biorbd_return], + ["q", "parameters"], + ["marker"], + ) + return casadi_fun + + @property + def nb_rigid_contacts(self) -> int: + """ + Returns the number of rigid contacts. + Example: + First contact with axis YZ + Second contact with axis Z + nb_rigid_contacts = 2 + """ + return self.model.nbRigidContacts() + + @property + def nb_contacts(self) -> int: + """ + Returns the number of contact index. + Example: + First contact with axis YZ + Second contact with axis Z + nb_contacts = 3 + """ + return self.model.nbContacts() + + def rigid_contact_index(self, contact_index) -> tuple: + """ + Returns the axis index of this specific rigid contact. + Example: + First contact with axis YZ + Second contact with axis Z + rigid_contact_index(0) = (1, 2) + """ + return self.model.rigidContacts()[contact_index].availableAxesIndices() + + def markers_velocities(self, reference_index=None) -> list[MX]: + if reference_index is None: + biorbd_return = [ + m.to_mx() + for m in self.model.markersVelocity( + GeneralizedCoordinates(self.q), + GeneralizedVelocity(self.qdot), + True, + ) + ] + + else: + biorbd_return = [] + homogeneous_matrix_transposed = self.homogeneous_matrices_in_global( + segment_index=reference_index, inverse=True + )( + GeneralizedCoordinates(self.q), + ) + for m in self.model.markersVelocity(GeneralizedCoordinates(self.q), GeneralizedVelocity(self.qdot)): + if m.applyRT(homogeneous_matrix_transposed) is None: + biorbd_return.append(m.to_mx()) + else: + biorbd_return.append(m.applyRT(homogeneous_matrix_transposed).to_mx()) + + casadi_fun = Function( + "markers_velocities", + [self.q, self.qdot, self.parameters], + [horzcat(*biorbd_return)], + ["q", "qdot", "parameters"], + ["markers_velocities"], + ) + return casadi_fun + + def marker_velocity(self, marker_index: int) -> list[MX]: + biorbd_return = self.model.markersVelocity( + GeneralizedCoordinates(self.q), + GeneralizedVelocity(self.qdot), + True, + )[marker_index].to_mx() + casadi_fun = Function( + "marker_velocity", + [self.q, self.qdot, self.parameters], + [biorbd_return], + ["q", "qdot", "parameters"], + ["marker_velocity"], + ) + return casadi_fun + + def markers_accelerations(self, reference_index=None) -> list[MX]: + if reference_index is None: + biorbd_return = [ + m.to_mx() + for m in self.model.markerAcceleration( + GeneralizedCoordinates(self.q), + GeneralizedVelocity(self.qdot), + GeneralizedAcceleration(self.qddot), + True, + ) + ] + + else: + biorbd_return = [] + homogeneous_matrix_transposed = self.homogeneous_matrices_in_global( + segment_index=reference_index, + inverse=True, + )( + GeneralizedCoordinates(self.q), + ) + for m in self.model.markersAcceleration( + GeneralizedCoordinates(self.q), + GeneralizedVelocity(self.qdot), + GeneralizedAcceleration(self.qddot), + ): + if m.applyRT(homogeneous_matrix_transposed) is None: + biorbd_return.append(m.to_mx()) + else: + biorbd_return.append(m.applyRT(homogeneous_matrix_transposed).to_mx()) + + casadi_fun = Function( + "markers_accelerations", + [self.q, self.qdot, self.qddot, self.parameters], + [horzcat(*biorbd_return)], + ["q", "qdot", "qddot", "parameters"], + ["markers_accelerations"], + ) + return casadi_fun + + def marker_acceleration(self, marker_index: int) -> list[MX]: + biorbd_return = self.model.markerAcceleration( + GeneralizedCoordinates(self.q), + GeneralizedVelocity(self.qdot), + GeneralizedAcceleration(self.qddot), + True, + )[marker_index].to_mx() + casadi_fun = Function( + "marker_acceleration", + [self.q, self.qdot, self.qddot, self.parameters], + [biorbd_return], + ["q", "qdot", "qddot", "parameters"], + ["marker_acceleration"], + ) + return casadi_fun + + def tau_max(self) -> tuple[MX, MX]: + self.model.closeActuator() + q_biorbd = GeneralizedCoordinates(self.q) + qdot_biorbd = GeneralizedVelocity(self.qdot) + torque_max, torque_min = self.model.torqueMax(q_biorbd, qdot_biorbd) + casadi_fun = Function( + "tau_max", + [self.q, self.qdot, self.parameters], + [torque_max.to_mx(), torque_min.to_mx()], + ["q", "qdot", "parameters"], + ["tau_max", "tau_min"], + ) + return casadi_fun + + def rigid_contact_acceleration(self, contact_index, contact_axis) -> Function: + q_biorbd = GeneralizedCoordinates(self.q) + qdot_biorbd = GeneralizedVelocity(self.qdot) + qddot_biorbd = GeneralizedAcceleration(self.qddot) + biorbd_return = self.model.rigidContactAcceleration( + q_biorbd, qdot_biorbd, qddot_biorbd, contact_index, True + ).to_mx()[contact_axis] + casadi_fun = Function( + "rigid_contact_acceleration", + [self.q, self.qdot, self.qddot, self.parameters], + [biorbd_return], + ["q", "qdot", "qddot", "parameters"], + ["rigid_contact_acceleration"], + ) + return casadi_fun + + def markers_jacobian(self) -> list[MX]: + biorbd_return = [m.to_mx() for m in self.model.markersJacobian(GeneralizedCoordinates(self.q))] + casadi_fun = Function( + "markers_jacobian", + [self.q, self.parameters], + biorbd_return, + ["q", "parameters"], + ["markers_jacobian"], + ) + return casadi_fun + + @property + def marker_names(self) -> tuple[str, ...]: + return tuple([s.to_string() for s in self.model.markerNames()]) + + def soft_contact_forces(self) -> Function: + q_biorbd = GeneralizedCoordinates(self.q) + qdot_biorbd = GeneralizedVelocity(self.qdot) + + biorbd_return = MX.zeros(self.nb_soft_contacts * 6, 1) + for i_sc in range(self.nb_soft_contacts): + soft_contact = self.soft_contact(i_sc) + biorbd_return[i_sc * 6 : (i_sc + 1) * 6, :] = ( + biorbd.SoftContactSphere(soft_contact).computeForceAtOrigin(self.model, q_biorbd, qdot_biorbd).to_mx() + ) + + casadi_fun = Function( + "soft_contact_forces", + [self.q, self.qdot, self.parameters], + [biorbd_return], + ["q", "qdot", "parameters"], + ["soft_contact_forces"], + ) + return casadi_fun + + def normalize_state_quaternions(self) -> Function: + + quat_idx = self.get_quaternion_idx() + biorbd_return = MX.zeros(self.nb_q) + biorbd_return[:] = self.q + + # Normalize quaternion, if needed + for j in range(self.nb_quaternions): + quaternion = vertcat( + self.q[quat_idx[j][3]], + self.q[quat_idx[j][0]], + self.q[quat_idx[j][1]], + self.q[quat_idx[j][2]], + ) + quaternion /= norm_fro(quaternion) + biorbd_return[quat_idx[j][0] : quat_idx[j][2] + 1] = quaternion[1:4] + biorbd_return[quat_idx[j][3]] = quaternion[0] + + casadi_fun = Function( + "normalize_state_quaternions", + [self.q], + [biorbd_return], + ["q"], + ["q_normalized"], + ) + return casadi_fun + + def get_quaternion_idx(self) -> list[list[int]]: + n_dof = 0 + quat_idx = [] + quat_number = 0 + for j in range(self.nb_segments): + if self.segments[j].isRotationAQuaternion(): + quat_idx.append([n_dof, n_dof + 1, n_dof + 2, self.nb_dof + quat_number]) + quat_number += 1 + n_dof += self.segments[j].nbDof() + return quat_idx + + def contact_forces(self) -> Function: + force = self.contact_forces_from_constrained_forward_dynamics()( + self.q, self.qdot, self.tau, self.external_forces, self.parameters + ) + casadi_fun = Function( + "contact_forces", + [self.q, self.qdot, self.tau, self.external_forces, self.parameters], + [force], + ["q", "qdot", "tau", "external_forces", "parameters"], + ["contact_forces"], + ) + return casadi_fun + + def passive_joint_torque(self) -> Function: + q_biorbd = GeneralizedCoordinates(self.q) + qdot_biorbd = GeneralizedVelocity(self.qdot) + biorbd_return = self.model.passiveJointTorque(q_biorbd, qdot_biorbd).to_mx() + casadi_fun = Function( + "passive_joint_torque", + [self.q, self.qdot, self.parameters], + [biorbd_return], + ["q", "qdot", "parameters"], + ["passive_joint_torque"], + ) + return casadi_fun + + def ligament_joint_torque(self) -> Function: + q_biorbd = GeneralizedCoordinates(self.q) + qdot_biorbd = GeneralizedVelocity(self.qdot) + biorbd_return = self.model.ligamentsJointTorque(q_biorbd, qdot_biorbd).to_mx() + casadi_fun = Function( + "ligament_joint_torque", + [self.q, self.qdot, self.parameters], + [biorbd_return], + ["q", "qdot", "parameters"], + ["ligament_joint_torque"], + ) + return casadi_fun + + def ranges_from_model(self, variable: str): + ranges = [] + for segment in self.segments: + if "_joints" in variable: + if segment.parent().to_string().lower() != "root": + variable = variable.replace("_joints", "") + ranges += self._add_range(variable, segment) + elif "_roots" in variable: + if segment.parent().to_string().lower() == "root": + variable = variable.replace("_roots", "") + ranges += self._add_range(variable, segment) + else: + ranges += self._add_range(variable, segment) + + return ranges + + @staticmethod + def _add_range(variable: str, segment: biorbd.Segment) -> list[biorbd.Range]: + """ + Get the range of a variable for a given segment + + Parameters + ---------- + variable: str + The variable to get the range for such as: + "q", "qdot", "qddot", "q_joint", "qdot_joint", "qddot_joint", "q_root", "qdot_root", "qddot_root" + segment: biorbd.Segment + The segment to get the range from + + Returns + ------- + list[biorbd.Range] + range min and max for the given variable for a given segment + """ + ranges_map = { + "q": [q_range for q_range in segment.QRanges()], + "qdot": [qdot_range for qdot_range in segment.QdotRanges()], + "qddot": [qddot_range for qddot_range in segment.QddotRanges()], + } + + segment_variable_range = ranges_map.get(variable, None) + if segment_variable_range is None: + RuntimeError("Wrong variable name") + + return segment_variable_range + + def _var_mapping( + self, + key: str, + range_for_mapping: int | list | tuple | range, + mapping: BiMapping = None, + ) -> dict: + return _var_mapping(key, range_for_mapping, mapping) + + def bounds_from_ranges(self, variables: str | list[str], mapping: BiMapping | BiMappingList = None) -> Bounds: + return bounds_from_ranges(self, variables, mapping) + + def lagrangian(self) -> Function: + q_biorbd = GeneralizedCoordinates(self.q) + qdot_biorbd = GeneralizedVelocity(self.qdot) + biorbd_return = self.model.Lagrangian(q_biorbd, qdot_biorbd).to_mx() + casadi_fun = Function( + "lagrangian", + [self.q, self.qdot], + [biorbd_return], + ["q", "qdot"], + ["lagrangian"], + ) + return casadi_fun + + def partitioned_forward_dynamics(self): + raise NotImplementedError("partitioned_forward_dynamics is not implemented for BiorbdModel") + + @staticmethod + def animate( + ocp, + solution, + show_now: bool = True, + show_tracked_markers: bool = False, + viewer: str = "pyorerun", + n_frames: int = 0, + **kwargs, + ): + if viewer == "bioviz": + from .viewer_bioviz import animate_with_bioviz_for_loop + + return animate_with_bioviz_for_loop(ocp, solution, show_now, show_tracked_markers, n_frames, **kwargs) + if viewer == "pyorerun": + from .viewer_pyorerun import animate_with_pyorerun + + return animate_with_pyorerun(ocp, solution, show_now, show_tracked_markers, **kwargs) From a44bf3085ce15de38f8bb7de43a03be2bb26ba6e Mon Sep 17 00:00:00 2001 From: Pariterre Date: Mon, 16 Dec 2024 12:28:14 -0500 Subject: [PATCH 04/19] Starting an example with torch --- .gitignore | 1 + .../custom_non_casadi_dynamics.py | 513 ++------ bioptim/models/torch/torch_model.py | 1064 +---------------- 3 files changed, 179 insertions(+), 1399 deletions(-) diff --git a/.gitignore b/.gitignore index 28ba13771..12a9d2595 100644 --- a/.gitignore +++ b/.gitignore @@ -122,6 +122,7 @@ _l4c_generated/ *.png +bioptim/examples/getting_started/_l4c_generated/ sandbox/ # Ignore all npy files in tests folder except those in tests/shard6 diff --git a/bioptim/examples/getting_started/custom_non_casadi_dynamics.py b/bioptim/examples/getting_started/custom_non_casadi_dynamics.py index 818b0cdbd..a32513ede 100644 --- a/bioptim/examples/getting_started/custom_non_casadi_dynamics.py +++ b/bioptim/examples/getting_started/custom_non_casadi_dynamics.py @@ -1,442 +1,191 @@ """ TODO: Explain what is this example about +TODO: All the documentation This example is similar to the getting_started/pendulum.py example, but the dynamics are computed using a non-casadi based model. This is useful when the dynamics are computed using a different library (e.g. TensorFlow, PyTorch, etc.) """ import biorbd -from bioptim import ( - OptimalControlProgram, - DynamicsFcn, - Objective, - ObjectiveFcn, - BoundsList, - OdeSolver, - OdeSolverBase, - PhaseDynamics, - ControlType, - InitialGuessList, - Dynamics, - CasadiFunctionInterface, - BiorbdModel, - PenaltyController, - Node, -) -from casadi import Function, jacobian, MX, DM - +from bioptim import OptimalControlProgram, DynamicsFcn, BoundsList, Dynamics +from bioptim.models.torch.torch_model import TorchModel import numpy as np - - -class CasadiFunctionInterfaceTest(CasadiFunctionInterface): - """ - This example implements a somewhat simple 5x1 function, with x and y inputs (x => 3x1; y => 4x1) of the form - f(x, y) = np.array( - [ - x[0] * y[1] + y[0] * y[0], - x[1] * x[1] + 2 * y[1], - x[0] * x[1] * x[2], - x[2] * x[1] + 2 * y[3] * y[2], - y[0] * y[1] * y[2] * y[3], - ] +import torch + + +class NeuralNetworkModel(torch.nn.Module): + def __init__( + self, + layer_node_count: tuple[int], + dropout_probability: float, + use_batch_norm: bool, + ): + super(NeuralNetworkModel, self).__init__() + activations = torch.nn.GELU() + + # Initialize the layers of the neural network + self._size_in = layer_node_count[0] + self._size_out = layer_node_count[-1] + first_and_hidden_layers_node_count = layer_node_count[:-1] + layers = torch.nn.ModuleList() + for i in range(len(first_and_hidden_layers_node_count) - 1): + layers.append( + torch.nn.Linear(first_and_hidden_layers_node_count[i], first_and_hidden_layers_node_count[i + 1]) + ) + if use_batch_norm: + torch.nn.BatchNorm1d(first_and_hidden_layers_node_count[i + 1]) + layers.append(activations) + layers.append(torch.nn.Dropout(dropout_probability)) + layers.append(torch.nn.Linear(first_and_hidden_layers_node_count[-1], layer_node_count[-1])) + + self._forward_model = torch.nn.Sequential(*layers) + self._forward_model.to(self.get_torch_device()) + + self._optimizer = torch.optim.Adam(self.parameters(), lr=1e-3) + self._loss_function = torch.nn.HuberLoss() + + # Put the model in evaluation mode + self.eval() + + @property + def size_in(self) -> int: + return self._size_in + + @property + def size_out(self) -> int: + return self._size_out + + def forward(self, x: torch.Tensor) -> torch.Tensor: + output = torch.Tensor(x.shape[0], self._forward_model[-1].out_features) + for i, data in enumerate(x): + output[i, :] = self._forward_model(data) + return output.to(self.get_torch_device()) + + @staticmethod + def get_torch_device() -> torch.device: + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + + def train_me(self, training_data: list[torch.Tensor], validation_data: list[torch.Tensor]): + # More details about scheduler in documentation + scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( + self._optimizer, mode="min", factor=0.1, patience=20, min_lr=1e-8 ) - It implements the equation (5x1) and the jacobians for the inputs x (5x3) and y (5x4). - """ - - def __init__(self, model, opts={}): - super(CasadiFunctionInterfaceTest, self).__init__("CasadiFunctionInterfaceTest", opts) - - def inputs_len(self) -> list[int]: - return [2, 2] + max_epochs = 10 + for _ in range(max_epochs): + self._perform_epoch_training(targets=training_data) + validation_loss = self._perform_epoch_training(targets=validation_data, only_compute=True) + print(f"Validation loss: {validation_loss}") + scheduler.step(validation_loss) # Adjust/reduce learning rate - def outputs_len(self) -> list[int]: - return [2] + def _perform_epoch_training( + self, + targets: list[torch.Tensor], + only_compute: bool = False, + ) -> tuple[float, float]: - def function(self, *args): - x, y = args - x = np.array(x)[:, 0] - y = np.array(y)[:, 0] - return np.array([x[0] * y[1] + x[0] * y[0] * y[0], x[1] * x[1] + 2 * y[1]]) + # Perform the predictions + if only_compute: + with torch.no_grad(): + all_predictions = self(targets[0]) + all_targets = targets[1] - def jacobians(self, *args): - x, y = args - x = np.array(x)[:, 0] - y = np.array(y)[:, 0] - jacobian_x = np.array([[y[1] + y[0] * y[0], 0, 0], [0, 2 * x[1], 0]]) - jacobian_y = np.array([[x[0] * 2 * y[0], x[0], 0, 0], [0, 2, 0, 0]]) - return [jacobian_x, jacobian_y] + else: + # Put the model in training mode + self.train() + # If it is training, we are updating the model with each prediction, we therefore need to do it in a loop + all_predictions = torch.tensor([]).to(self.get_torch_device()) + all_targets = torch.tensor([]).to(self.get_torch_device()) + for input, target in zip(*targets): + self._optimizer.zero_grad() -class ForwardDynamicsInterface(CasadiFunctionInterface): - def __init__(self, model: BiorbdModel, opts={}): - self.non_casadi_model = biorbd.Model(model.path) - super(ForwardDynamicsInterface, self).__init__("ForwardDynamicsInterface", opts) + # Get the predictions and targets + output = self(input[None, :]) - def inputs_len(self) -> list[int]: - return [1] + # Do some machine learning shenanigans + current_loss = self._loss_function.forward(output, target[None, :]) + current_loss.backward() # Backpropagation + self._optimizer.step() # Updating weights - def outputs_len(self) -> list[int]: - return [1] + # Populate the return values + all_predictions = torch.cat((all_predictions, output)) + all_targets = torch.cat((all_targets, target[None, :])) - def function(self, *args): - return [args[0]] # self.non_casadi_model.ForwardDynamics(*self.mx_in()[:3]) + # Put back the model in evaluation mode + self.eval() - def jacobians(self, *args): - return [0, 0, DM(1), 0, 0] - - -def custom_func_track_markers(controller: PenaltyController) -> MX: - return controller.model.custom_interface(controller.states["q"].cx, controller.controls["tau"].cx) + # Calculation of mean distance and error % + epoch_accuracy = (all_predictions - all_targets).abs().mean().item() + return epoch_accuracy def prepare_ocp( - biorbd_model_path: str, + model: torch.nn.Module, final_time: float, n_shooting: int, - ode_solver: OdeSolverBase = OdeSolver.RK4(), - use_sx: bool = False, - n_threads: int = 1, - phase_dynamics: PhaseDynamics = PhaseDynamics.SHARED_DURING_THE_PHASE, - expand_dynamics: bool = True, - control_type: ControlType = ControlType.CONSTANT, ) -> OptimalControlProgram: - """ - The initialization of an ocp - - Parameters - ---------- - biorbd_model_path: str - The path to the biorbd model - final_time: float - The time in second required to perform the task - n_shooting: int - The number of shooting points to define int the direct multiple shooting program - ode_solver: OdeSolverBase = OdeSolver.RK4() - Which type of OdeSolver to use - use_sx: bool - If the SX variable should be used instead of MX (can be extensive on RAM) - n_threads: int - The number of threads to use in the paralleling (1 = no parallel computing) - phase_dynamics: PhaseDynamics - If the dynamics equation within a phase is unique or changes at each node. - PhaseDynamics.SHARED_DURING_THE_PHASE is much faster, but lacks the capability to have changing dynamics within - a phase. A good example of when PhaseDynamics.ONE_PER_NODE should be used is when different external forces - are applied at each node - expand_dynamics: bool - If the dynamics function should be expanded. Please note, this will solve the problem faster, but will slow down - the declaration of the OCP, so it is a trade-off. Also depending on the solver, it may or may not work - (for instance IRK is not compatible with expanded dynamics) - control_type: ControlType - The type of the controls - - Returns - ------- - The OptimalControlProgram ready to be solved - """ - - bio_model = BiorbdModel(biorbd_model_path) - bio_model.custom_interface = CasadiFunctionInterfaceTest(bio_model) - - # Add objective functions - objective_functions = Objective(custom_func_track_markers, custom_type=ObjectiveFcn.Mayer, node=Node.START) # Dynamics - dynamics = Dynamics(DynamicsFcn.TORQUE_DRIVEN, expand_dynamics=expand_dynamics, phase_dynamics=phase_dynamics) + dynamics = Dynamics(DynamicsFcn.TORQUE_DRIVEN) + torch_model = TorchModel(torch_model=model) # Path bounds x_bounds = BoundsList() - x_bounds["q"] = bio_model.bounds_from_ranges("q") + x_bounds["q"] = [-3.14 * 1.5] * torch_model.nb_q, [3.14 * 1.5] * torch_model.nb_q x_bounds["q"][:, [0, -1]] = 0 # Start and end at 0... x_bounds["q"][1, -1] = 3.14 # ...but end with pendulum 180 degrees rotated - x_bounds["qdot"] = bio_model.bounds_from_ranges("qdot") + x_bounds["qdot"] = [-3.14 * 10.0] * torch_model.nb_qdot, [3.14 * 10.0] * torch_model.nb_qdot x_bounds["qdot"][:, [0, -1]] = 0 # Start and end without any velocity - # Initial guess (optional since it is 0, we show how to initialize anyway) - x_init = InitialGuessList() - x_init["q"] = [0] * bio_model.nb_q - x_init["qdot"] = [0] * bio_model.nb_qdot - # Define control path bounds - n_tau = bio_model.nb_tau u_bounds = BoundsList() - u_bounds["tau"] = [-100] * n_tau, [100] * n_tau # Limit the strength of the pendulum to (-100 to 100)... + u_bounds["tau"] = [-100] * torch_model.nb_tau, [100] * torch_model.nb_tau u_bounds["tau"][1, :] = 0 # ...but remove the capability to actively rotate - # Initial guess (optional since it is 0, we show how to initialize anyway) - u_init = InitialGuessList() - u_init["tau"] = [0] * n_tau - return OptimalControlProgram( - bio_model, + torch_model, dynamics, n_shooting, final_time, - x_init=x_init, - u_init=u_init, x_bounds=x_bounds, u_bounds=u_bounds, - objective_functions=objective_functions, - ode_solver=ode_solver, - use_sx=use_sx, - n_threads=n_threads, - control_type=control_type, + use_sx=True, ) +def generate_dataset(biorbd_model: biorbd.Model, data_point_count: int) -> list[torch.Tensor]: + q = torch.rand(data_point_count, biorbd_model.nbQ()) + qdot = torch.rand(data_point_count, biorbd_model.nbQdot()) + tau = torch.rand(data_point_count, biorbd_model.nbGeneralizedTorque()) + + qddot = torch.zeros(data_point_count, biorbd_model.nbQddot()) + for i in range(data_point_count): + qddot[i, :] = torch.tensor( + biorbd_model.ForwardDynamics(np.array(q[i, :]), np.array(qdot[i, :]), np.array(tau[i, :])).to_array() + ) + + return [torch.cat((q, qdot, tau), dim=1), qddot] + + def main(): - """ - If pendulum is run as a script, it will perform the optimization and animates it - """ + # --- Prepare a predictive model --- # + biorbd_model = biorbd.Model("models/pendulum.bioMod") + training_data = generate_dataset(biorbd_model, data_point_count=1000) + validation_data = generate_dataset(biorbd_model, data_point_count=100) + + model = NeuralNetworkModel(layer_node_count=(6, 10, 10, 2), dropout_probability=0.2, use_batch_norm=True) + model.train_me(training_data, validation_data) - # --- Prepare the ocp --- # - ocp = prepare_ocp(biorbd_model_path="models/pendulum.bioMod", final_time=1, n_shooting=400, n_threads=2) + ocp = prepare_ocp(model=model, final_time=1, n_shooting=40) # --- Solve the ocp --- # sol = ocp.solve() # --- Show the results graph --- # # sol.print_cost() - sol.graphs(show_bounds=True, save_name="results.png") + sol.graphs(show_bounds=True) if __name__ == "__main__": main() - - -######## OCP FAST ######## -# from casadi import * - -# T = 10.0 # Time horizon -# N = 20 # number of control intervals - -# # Declare model variables -# x1 = MX.sym("x1") -# x2 = MX.sym("x2") -# x = vertcat(x1, x2) -# u = MX.sym("u") - -# # Model equations -# xdot = vertcat((1 - x2**2) * x1 - x2 + u, x1) - - -# # Formulate discrete time dynamics -# if False: -# # CVODES from the SUNDIALS suite -# dae = {"x": x, "p": u, "ode": xdot} -# opts = {"tf": T / N} -# F = integrator("F", "cvodes", dae, opts) -# else: -# # Fixed step Runge-Kutta 4 integrator -# M = 4 # RK4 steps per interval -# DT = T / N / M -# f = Function("f", [x, u], [xdot]) -# X0 = MX.sym("X0", 2) -# U = MX.sym("U") -# X = X0 -# Q = 0 -# for j in range(M): -# k1 = f(X, U) -# k2 = f(X + DT / 2 * k1, U) -# k3 = f(X + DT / 2 * k2, U) -# k4 = f(X + DT * k3, U) -# X = X + DT / 6 * (k1 + 2 * k2 + 2 * k3 + k4) -# F = Function("F", [X0, U], [X], ["x0", "p"], ["xf"]) - -# # Start with an empty NLP -# w = [] -# w0 = [] -# lbw = [] -# ubw = [] -# g = [] -# lbg = [] -# ubg = [] - -# # "Lift" initial conditions -# Xk = MX.sym("X0", 2) -# w += [Xk] -# lbw += [0, 1] -# ubw += [0, 1] -# w0 += [0, 1] - -# # Formulate the NLP -# for k in range(N): -# # New NLP variable for the control -# Uk = MX.sym("U_" + str(k)) -# w += [Uk] -# lbw += [-1] -# ubw += [1] -# w0 += [0] - -# # Integrate till the end of the interval -# Fk = F(x0=Xk, p=Uk) -# Xk_end = Fk["xf"] - -# # New NLP variable for state at end of interval -# Xk = MX.sym("X_" + str(k + 1), 2) -# w += [Xk] -# lbw += [-0.25, -inf] -# ubw += [inf, inf] -# w0 += [0, 0] - -# # Add equality constraint -# g += [Xk_end - Xk] -# lbg += [0, 0] -# ubg += [0, 0] - -# nd = N + 1 - -# import gpflow -# import time - -# from tensorflow_casadi import TensorFlowEvaluator - - -# class GPR(TensorFlowEvaluator): -# def __init__(self, session, opts={}): -# X = tf.compat.v1.placeholder(shape=(1, nd), dtype=np.float64) -# mean = tf.reshape(tf.reduce_mean(X), (1, 1)) -# TensorFlowEvaluator.__init__(self, [X], [mean], session, opts) -# self.counter = 0 -# self.time = 0 - -# def eval(self, arg): -# self.counter += 1 -# t0 = time.time() -# ret = TensorFlowEvaluator.eval(self, arg) -# self.time += time.time() - t0 -# return [ret] - - -# import tensorflow as tf - -# with tf.compat.v1.Session() as session: -# GPR = GPR(session) - -# w = vertcat(*w) - -# # Create an NLP solver -# prob = {"f": sum1(GPR(w[0::3])), "x": w, "g": vertcat(*g)} -# options = {"ipopt": {"hessian_approximation": "limited-memory"}} -# solver = nlpsol("solver", "ipopt", prob, options) - -# # Solve the NLP -# sol = solver(x0=w0, lbx=lbw, ubx=ubw, lbg=lbg, ubg=ubg) - -# print("Ncalls", GPR.counter) -# print("Total time [s]", GPR.time) -# w_opt = sol["x"].full().flatten() - -# # Plot the solution -# x1_opt = w_opt[0::3] -# x2_opt = w_opt[1::3] -# u_opt = w_opt[2::3] - -# tgrid = [T / N * k for k in range(N + 1)] -# import matplotlib.pyplot as plt - -# plt.figure(1) -# plt.clf() -# plt.plot(tgrid, x1_opt, "--") -# plt.plot(tgrid, x2_opt, "-") -# plt.step(tgrid, vertcat(DM.nan(1), u_opt), "-.") -# plt.xlabel("t") -# plt.legend(["x1", "x2", "u"]) -# plt.grid() -# plt.show() - - -# -# -# -######### TENSORFLOW CASADI ######### -# import casadi -# import tensorflow as tf - - -# class TensorFlowEvaluator(casadi.Callback): -# def __init__(self, t_in, t_out, session, opts={}): -# """ -# t_in: list of inputs (tensorflow placeholders) -# t_out: list of outputs (tensors dependeant on those placeholders) -# session: a tensorflow session -# """ -# casadi.Callback.__init__(self) -# assert isinstance(t_in, list) -# self.t_in = t_in -# assert isinstance(t_out, list) -# self.t_out = t_out -# self.construct("TensorFlowEvaluator", opts) -# self.session = session -# self.refs = [] - -# def get_n_in(self): -# return len(self.t_in) - -# def get_n_out(self): -# return len(self.t_out) - -# def get_sparsity_in(self, i): -# return casadi.Sparsity.dense(*self.t_in[i].get_shape().as_list()) - -# def get_sparsity_out(self, i): -# return casadi.Sparsity.dense(*self.t_out[i].get_shape().as_list()) - -# def eval(self, arg): -# # Associate each tensorflow input with the numerical argument passed by CasADi -# d = dict((v, arg[i].toarray()) for i, v in enumerate(self.t_in)) -# # Evaluate the tensorflow expressions -# ret = self.session.run(self.t_out, feed_dict=d) -# return ret - -# # Vanilla tensorflow offers just the reverse mode AD -# def has_reverse(self, nadj): -# return nadj == 1 - -# def get_reverse(self, nadj, name, inames, onames, opts): -# # Construct tensorflow placeholders for the reverse seeds -# adj_seed = [ -# tf.compat.v1.placeholder(shape=self.sparsity_out(i).shape, dtype=tf.float64) for i in range(self.n_out()) -# ] -# # Construct the reverse tensorflow graph through 'gradients' -# grad = tf.gradients(self.t_out, self.t_in, grad_ys=adj_seed) -# # Create another TensorFlowEvaluator object -# callback = TensorFlowEvaluator(self.t_in + adj_seed, grad, self.session) -# # Make sure you keep a reference to it -# self.refs.append(callback) - -# # Package it in the nominal_in+nominal_out+adj_seed form that CasADi expects -# nominal_in = self.mx_in() -# nominal_out = self.mx_out() -# adj_seed = self.mx_out() -# return casadi.Function( -# name, nominal_in + nominal_out + adj_seed, callback.call(nominal_in + adj_seed), inames, onames -# ) - - -# if __name__ == "__main__": -# from casadi import * - -# a = tf.compat.v1.placeholder(shape=(2, 2), dtype=tf.float64) -# b = tf.compat.v1.placeholder(shape=(2, 1), dtype=tf.float64) - -# y = tf.matmul(tf.sin(a), b) - -# with tf.compat.v1.Session() as session: -# f_tf = TensorFlowEvaluator([a, b], [y], session) - -# a = MX.sym("a", 2, 2) -# b = MX.sym("a", 2, 1) -# y = f_tf(a, b) -# yref = mtimes(sin(a), b) - -# f = Function("f", [a, b], [y]) -# fref = Function("f", [a, b], [yref]) - -# print(f(DM([[1, 2], [3, 4]]), DM([[1], [3]]))) -# print(fref(DM([[1, 2], [3, 4]]), DM([[1], [3]]))) - -# f = Function("f", [a, b], [jacobian(y, a)]) -# fref = Function("f", [a, b], [jacobian(yref, a)]) -# print(f(DM([[1, 2], [3, 4]]), DM([[1], [3]]))) -# print(fref(DM([[1, 2], [3, 4]]), DM([[1], [3]]))) diff --git a/bioptim/models/torch/torch_model.py b/bioptim/models/torch/torch_model.py index 0180bd35a..abdfcf28b 100644 --- a/bioptim/models/torch/torch_model.py +++ b/bioptim/models/torch/torch_model.py @@ -1,20 +1,24 @@ from typing import Callable from casadi import SX, MX, vertcat, horzcat, norm_fro, Function +import l4casadi as l4c import numpy as np import torch -""" -INSTALLATION: -First, make sure pytorch is installed - - pip install torch>=2.0 --index-url https://download.pytorch.org/whl/cpu/torch_stable.html +# """ +# INSTALLATION: +# First, make sure pytorch is installed -Then, install l4casadi as the interface between CasADi and PyTorch +# pip install torch>=2.0 --index-url https://download.pytorch.org/whl/cpu/torch_stable.html +# setuptools>=68.1 +# scikit-build>=0.17 +# cmake>=3.27 +# ninja>=1.11 +# Then, install l4casadi as the interface between CasADi and PyTorch - pip install l4casadi --no-build-isolation +# pip install l4casadi --no-build-isolation -""" +# """ class TorchModel: @@ -22,1050 +26,76 @@ class TorchModel: This class wraps a pytorch model and allows the user to call some useful functions on it. """ - def __init__(self, model: str | torch.nn.Module): - if not isinstance(bio_model, str) and not isinstance(bio_model, biorbd.Model): - raise ValueError("The model should be of type 'str' or 'biorbd.Model'") + def __init__(self, torch_model: torch.nn.Module): + self._dynamic_model = l4c.L4CasADi(torch_model, device="cpu") # device='cuda' for GPU - self.model = biorbd.Model(bio_model) if isinstance(bio_model, str) else bio_model - if parameters is not None: - for param_key in parameters: - parameters[param_key].apply_parameter(self) - self._friction_coefficients = friction_coefficients - - self.external_force_set = ( - self._set_external_force_set(external_force_set) if external_force_set is not None else None - ) + self._nb_dof = torch_model.size_in // 3 self._symbolic_variables() - self.biorbd_external_forces_set = self._dispatch_forces() if external_force_set else None - - # TODO: remove mx (the MX parameters should be created inside the BiorbdModel) - self.parameters = parameters.mx if parameters else MX() def _symbolic_variables(self): """Declaration of MX variables of the right shape for the creation of CasADi Functions""" - self.q = MX.sym("q_mx", self.nb_q, 1) - self.qdot = MX.sym("qdot_mx", self.nb_qdot, 1) - self.qddot = MX.sym("qddot_mx", self.nb_qddot, 1) - self.qddot_joints = MX.sym("qddot_joints_mx", self.nb_qddot - self.nb_root, 1) - self.tau = MX.sym("tau_mx", self.nb_tau, 1) - self.muscle = MX.sym("muscle_mx", self.nb_muscles, 1) - self.activations = MX.sym("activations_mx", self.nb_muscles, 1) - self.external_forces = MX.sym( - "external_forces_mx", - self.external_force_set.nb_external_forces_components if self.external_force_set else 0, - 1, - ) - - def _set_external_force_set(self, external_force_set: ExternalForceSetTimeSeries): - """ - It checks the external forces and binds them to the model. - """ - external_force_set._check_segment_names(tuple([s.name().to_string() for s in self.model.segments()])) - external_force_set._check_all_string_points_of_application(self.marker_names) - external_force_set._bind() - - return external_force_set + self.q = MX.sym("q_mx", self.nb_dof, 1) + self.qdot = MX.sym("qdot_mx", self.nb_dof, 1) + self.tau = MX.sym("tau_mx", self.nb_dof, 1) + self.external_forces = MX.sym("external_forces_mx", 0, 1) + self.parameters = MX.sym("parameters_mx", 0, 1) @property def name(self) -> str: # parse the path and split to get the .bioMod name - return self.model.path().absolutePath().to_string().split("/")[-1] - - @property - def path(self) -> str: - return self.model.path().relativePath().to_string() - - def copy(self): - return BiorbdModel(self.path) - - def serialize(self) -> tuple[Callable, dict]: - return BiorbdModel, dict(bio_model=self.path) - - @property - def friction_coefficients(self) -> MX | SX | np.ndarray: - return self._friction_coefficients - - def set_friction_coefficients(self, new_friction_coefficients) -> None: - if np.any(new_friction_coefficients < 0): - raise ValueError("Friction coefficients must be positive") - return self._friction_coefficients - - @property - def gravity(self) -> Function: - """ - Returns the gravity of the model. - Since the gravity is self-defined in the model, you need to provide the type of the output when calling the function like this: - model.gravity()(MX() / SX()) - """ - biorbd_return = self.model.getGravity().to_mx() - casadi_fun = Function( - "gravity", - [self.parameters], - [biorbd_return], - ["parameters"], - ["gravity"], - ) - return casadi_fun - - def set_gravity(self, new_gravity) -> None: - self.model.setGravity(new_gravity) - return - - @property - def nb_tau(self) -> int: - return self.model.nbGeneralizedTorque() - - @property - def nb_segments(self) -> int: - return self.model.nbSegment() - - def segment_index(self, name) -> int: - return biorbd.segment_index(self.model, name) + return "forward_dynamics_torch_model" @property - def nb_quaternions(self) -> int: - return self.model.nbQuat() + def name_dof(self) -> list[str]: + return [f"q_{i}" for i in range(self.nb_dof)] @property def nb_dof(self) -> int: - return self.model.nbDof() + return self._nb_dof @property def nb_q(self) -> int: - return self.model.nbQ() + return self.nb_dof @property def nb_qdot(self) -> int: - return self.model.nbQdot() - - @property - def nb_qddot(self) -> int: - return self.model.nbQddot() - - @property - def nb_root(self) -> int: - return self.model.nbRoot() - - @property - def segments(self) -> tuple[biorbd.Segment]: - return self.model.segments() - - def rotation_matrix_to_euler_angles(self, sequence: str) -> Function: - """ - Returns the rotation matrix to euler angles function. - """ - r = MX.sym("r_mx", 3, 3) - r_matrix = biorbd.Rotation(r[0, 0], r[0, 1], r[0, 2], r[1, 0], r[1, 1], r[1, 2], r[2, 0], r[2, 1], r[2, 2]) - biorbd_return = biorbd.Rotation.toEulerAngles(r_matrix, sequence).to_mx() - casadi_fun = Function( - "rotation_matrix_to_euler_angles", - [r], - [biorbd_return], - ["Rotation matrix"], - ["Euler angles"], - ) - return casadi_fun - - def homogeneous_matrices_in_global(self, segment_index: int, inverse=False) -> Function: - """ - Returns the roto-translation matrix of the segment in the global reference frame. - """ - q_biorbd = GeneralizedCoordinates(self.q) - jcs = self.model.globalJCS(q_biorbd, segment_index) - biorbd_return = jcs.transpose().to_mx() if inverse else jcs.to_mx() - casadi_fun = Function( - "homogeneous_matrices_in_global", - [self.q, self.parameters], - [biorbd_return], - ["q", "parameters"], - ["Joint coordinate system RT matrix in global"], - ) - return casadi_fun - - def homogeneous_matrices_in_child(self, segment_id) -> Function: - """ - Returns the roto-translation matrix of the segment in the child reference frame. - Since the homogeneous matrix is self-defined in the model, you need to provide the type of the output when calling the function like this: - model.homogeneous_matrices_in_child(segment_id)(MX() / SX()) - """ - biorbd_return = self.model.localJCS(segment_id).to_mx() - casadi_fun = Function( - "homogeneous_matrices_in_child", - [self.parameters], - [biorbd_return], - ["parameters"], - ["Joint coordinate system RT matrix in local"], - ) - return casadi_fun - - @property - def mass(self) -> Function: - """ - Returns the mass of the model. - Since the mass is self-defined in the model, you need to provide the type of the output when calling the function like this: - model.mass()(MX() / SX()) - """ - biorbd_return = self.model.mass().to_mx() - casadi_fun = Function( - "mass", - [self.parameters], - [biorbd_return], - ["parameters"], - ["mass"], - ) - return casadi_fun - - def rt(self, rt_index) -> Function: - q_biorbd = GeneralizedCoordinates(self.q) - biorbd_return = self.model.RT(q_biorbd, rt_index).to_mx() - casadi_fun = Function( - "rt", - [self.q, self.parameters], - [biorbd_return], - ["q", "parameters"], - ["RT matrix"], - ) - return casadi_fun - - def center_of_mass(self) -> Function: - q_biorbd = GeneralizedCoordinates(self.q) - biorbd_return = self.model.CoM(q_biorbd, True).to_mx() - casadi_fun = Function( - "center_of_mass", - [self.q, self.parameters], - [biorbd_return], - ["q", "parameters"], - ["Center of mass"], - ) - return casadi_fun - - def center_of_mass_velocity(self) -> Function: - q_biorbd = GeneralizedCoordinates(self.q) - qdot_biorbd = GeneralizedVelocity(self.qdot) - biorbd_return = self.model.CoMdot(q_biorbd, qdot_biorbd, True).to_mx() - casadi_fun = Function( - "center_of_mass_velocity", - [self.q, self.qdot, self.parameters], - [biorbd_return], - ["q", "qdot", "parameters"], - ["Center of mass velocity"], - ) - return casadi_fun - - def center_of_mass_acceleration(self) -> Function: - q_biorbd = GeneralizedCoordinates(self.q) - qdot_biorbd = GeneralizedVelocity(self.qdot) - qddot_biorbd = GeneralizedAcceleration(self.qddot) - biorbd_return = self.model.CoMddot(q_biorbd, qdot_biorbd, qddot_biorbd, True).to_mx() - casadi_fun = Function( - "center_of_mass_acceleration", - [self.q, self.qdot, self.qddot, self.parameters], - [biorbd_return], - ["q", "qdot", "qddot", "parameters"], - ["Center of mass acceleration"], - ) - return casadi_fun - - def body_rotation_rate(self) -> Function: - q_biorbd = GeneralizedCoordinates(self.q) - qdot_biorbd = GeneralizedVelocity(self.qdot) - biorbd_return = self.model.bodyAngularVelocity(q_biorbd, qdot_biorbd, True).to_mx() - casadi_fun = Function( - "body_rotation_rate", - [self.q, self.qdot, self.parameters], - [biorbd_return], - ["q", "qdot", "parameters"], - ["Body rotation rate"], - ) - return casadi_fun - - def mass_matrix(self) -> Function: - q_biorbd = GeneralizedCoordinates(self.q) - biorbd_return = self.model.massMatrix(q_biorbd).to_mx() - casadi_fun = Function( - "mass_matrix", - [self.q, self.parameters], - [biorbd_return], - ["q", "parameters"], - ["Mass matrix"], - ) - return casadi_fun - - def non_linear_effects(self) -> Function: - q_biorbd = GeneralizedCoordinates(self.q) - qdot_biorbd = GeneralizedVelocity(self.qdot) - biorbd_return = self.model.NonLinearEffect(q_biorbd, qdot_biorbd).to_mx() - casadi_fun = Function( - "non_linear_effects", - [self.q, self.qdot, self.parameters], - [biorbd_return], - ["q", "qdot", "parameters"], - ["Non linear effects"], - ) - return casadi_fun - - def angular_momentum(self) -> Function: - q_biorbd = GeneralizedCoordinates(self.q) - qdot_biorbd = GeneralizedVelocity(self.qdot) - biorbd_return = self.model.angularMomentum(q_biorbd, qdot_biorbd, True).to_mx() - casadi_fun = Function( - "angular_momentum", - [self.q, self.qdot, self.parameters], - [biorbd_return], - ["q", "qdot", "parameters"], - ["Angular momentum"], - ) - return casadi_fun - - def reshape_qdot(self, k_stab=1) -> Function: - biorbd_return = self.model.computeQdot( - GeneralizedCoordinates(self.q), - GeneralizedCoordinates(self.qdot), # mistake in biorbd - k_stab, - ).to_mx() - casadi_fun = Function( - "reshape_qdot", - [self.q, self.qdot, self.parameters], - [biorbd_return], - ["q", "qdot", "parameters"], - ["Reshaped qdot"], - ) - return casadi_fun - - def segment_angular_velocity(self, idx) -> Function: - """ - Returns the angular velocity of the segment in the global reference frame. - """ - q_biorbd = GeneralizedCoordinates(self.q) - qdot_biorbd = GeneralizedVelocity(self.qdot) - biorbd_return = self.model.segmentAngularVelocity(q_biorbd, qdot_biorbd, idx, True).to_mx() - casadi_fun = Function( - "segment_angular_velocity", - [self.q, self.qdot, self.parameters], - [biorbd_return], - ["q", "qdot", "parameters"], - ["Segment angular velocity"], - ) - return casadi_fun - - def segment_orientation(self, idx: int, sequence: str = "xyz") -> Function: - """ - Returns the angular position of the segment in the global reference frame. - """ - q_biorbd = GeneralizedCoordinates(self.q) - rotation_matrix = self.homogeneous_matrices_in_global(idx)(q_biorbd, self.parameters)[:3, :3] - biorbd_return = self.rotation_matrix_to_euler_angles(sequence=sequence)(rotation_matrix) - casadi_fun = Function( - "segment_orientation", - [self.q, self.parameters], - [biorbd_return], - ["q", "parameters"], - ["Segment orientation"], - ) - return casadi_fun - - @property - def name_dof(self) -> tuple[str, ...]: - return tuple(s.to_string() for s in self.model.nameDof()) - - @property - def contact_names(self) -> tuple[str, ...]: - return tuple(s.to_string() for s in self.model.contactNames()) - - @property - def nb_soft_contacts(self) -> int: - return self.model.nbSoftContacts() - - @property - def soft_contact_names(self) -> tuple[str, ...]: - return self.model.softContactNames() - - def soft_contact(self, soft_contact_index, *args): - return self.model.softContact(soft_contact_index, *args) + return self.nb_dof @property - def muscle_names(self) -> tuple[str, ...]: - return tuple(s.to_string() for s in self.model.muscleNames()) - - @property - def nb_muscles(self) -> int: - return self.model.nbMuscles() - - def torque(self) -> Function: - """ - Returns the torque from the torque_activations. - Note that tau_activation should be between 0 and 1. - """ - q_biorbd = GeneralizedCoordinates(self.q) - qdot_biorbd = GeneralizedVelocity(self.qdot) - tau_activations_biorbd = self.tau - biorbd_return = self.model.torque(tau_activations_biorbd, q_biorbd, qdot_biorbd).to_mx() - casadi_fun = Function( - "torque_activation", - [self.tau, self.q, self.qdot, self.parameters], - [biorbd_return], - ["tau", "q", "qdot", "parameters"], - ["Torque from tau activations"], - ) - return casadi_fun - - def forward_dynamics_free_floating_base(self) -> Function: - q_biorbd = GeneralizedCoordinates(self.q) - qdot_biorbd = GeneralizedVelocity(self.qdot) - qddot_joints_biorbd = GeneralizedAcceleration(self.qddot_joints) - biorbd_return = self.model.ForwardDynamicsFreeFloatingBase(q_biorbd, qdot_biorbd, qddot_joints_biorbd).to_mx() - casadi_fun = Function( - "forward_dynamics_free_floating_base", - [self.q, self.qdot, self.qddot_joints, self.parameters], - [biorbd_return], - ["q", "qdot", "qddot_joints", "parameters"], - ["qddot_root and qddot_joints"], - ) - return casadi_fun - - @staticmethod - def reorder_qddot_root_joints(qddot_root, qddot_joints) -> MX | SX: - return vertcat(qddot_root, qddot_joints) - - def _dispatch_forces(self) -> biorbd.ExternalForceSet: - """Dispatch the symbolic MX into the biorbd external forces object""" - biorbd_external_forces = self.model.externalForceSet() - - # "type of external force": (function to call, number of force components) - force_mapping = { - "in_global": (_add_global_force, 6), - "torque_in_global": (_add_torque_global, 3), - "translational_in_global": (_add_translational_global, 3), - "in_local": (_add_local_force, 6), - "torque_in_local": (_add_torque_local, 3), - } - - symbolic_counter = 0 - for force_type, val in force_mapping.items(): - add_force_func, num_force_components = val - symbolic_counter = self._dispatch_forces_of_type( - force_type, add_force_func, num_force_components, symbolic_counter, biorbd_external_forces - ) - - return biorbd_external_forces - - def _dispatch_forces_of_type( - self, - force_type: str, - add_force_func: "Callable", - num_force_components: int, - symbolic_counter: int, - biorbd_external_forces: "biorbd.ExternalForces", - ) -> int: - """ - Helper method to dispatch forces of a specific external forces. - - Parameters - ---------- - force_type: str - The type of external force to dispatch among in_global, torque_in_global, translational_in_global, in_local, torque_in_local. - add_force_func: Callable - The function to call to add the force to the biorbd external forces object. - num_force_components: int - The number of force components for the given type - symbolic_counter: int - The current symbolic counter to slice the whole external_forces mx. - biorbd_external_forces: biorbd.ExternalForces - The biorbd external forces object to add the forces to. - - Returns - ------- - int - The updated symbolic counter. - """ - for segment, forces_on_segment in getattr(self.external_force_set, force_type).items(): - for force in forces_on_segment: - force_slicer = slice(symbolic_counter, symbolic_counter + num_force_components) - - point_of_application_mx = self._get_point_of_application(force, force_slicer.stop) - - add_force_func( - biorbd_external_forces, segment, self.external_forces[force_slicer], point_of_application_mx - ) - symbolic_counter = force_slicer.stop + ( - 3 if isinstance(force["point_of_application"], np.ndarray) else 0 - ) - - return symbolic_counter - - def _get_point_of_application(self, force, stop_index) -> biorbd.NodeSegment | np.ndarray | None: - """ - Determine the point of application mx slice based on its type. Only sliced if an array is stored - - Parameters - ---------- - force : dict - The force dictionary with details on the point of application. - stop_index : int - Index position in MX where the point of application components start. - - Returns - ------- - biorbd.NodeSegment | np.ndarray | None - Returns a slice of MX, a marker node, or None if no point of application is defined. - """ - if isinstance(force["point_of_application"], np.ndarray): - return self.external_forces[slice(stop_index, stop_index + 3)] - elif isinstance(force["point_of_application"], str): - return self.model.marker(self.marker_index(force["point_of_application"])) - return None + def nb_tau(self) -> int: + return self.nb_dof def forward_dynamics(self, with_contact: bool = False) -> Function: - - q_biorbd = GeneralizedCoordinates(self.q) - qdot_biorbd = GeneralizedVelocity(self.qdot) - tau_biorbd = GeneralizedTorque(self.tau) - - if with_contact: - if self.external_force_set is None: - biorbd_return = self.model.ForwardDynamicsConstraintsDirect(q_biorbd, qdot_biorbd, tau_biorbd).to_mx() - else: - biorbd_return = self.model.ForwardDynamicsConstraintsDirect( - q_biorbd, qdot_biorbd, tau_biorbd, self.biorbd_external_forces_set - ).to_mx() - casadi_fun = Function( - "constrained_forward_dynamics", - [self.q, self.qdot, self.tau, self.external_forces, self.parameters], - [biorbd_return], - ["q", "qdot", "tau", "external_forces", "parameters"], - ["qddot"], - ) - else: - if self.external_force_set is None: - biorbd_return = self.model.ForwardDynamics(q_biorbd, qdot_biorbd, tau_biorbd).to_mx() - else: - biorbd_return = self.model.ForwardDynamics( - q_biorbd, qdot_biorbd, tau_biorbd, self.biorbd_external_forces_set - ).to_mx() - casadi_fun = Function( - "forward_dynamics", - [self.q, self.qdot, self.tau, self.external_forces, self.parameters], - [biorbd_return], - ["q", "qdot", "tau", "external_forces", "parameters"], - ["qddot"], - ) - return casadi_fun - - def inverse_dynamics(self, with_contact: bool = False) -> Function: - - if with_contact: - raise NotImplementedError("Inverse dynamics with contact is not implemented yet") - - q_biorbd = GeneralizedCoordinates(self.q) - qdot_biorbd = GeneralizedVelocity(self.qdot) - qddot_biorbd = GeneralizedAcceleration(self.qddot) - if self.external_force_set is None: - biorbd_return = self.model.InverseDynamics(q_biorbd, qdot_biorbd, qddot_biorbd).to_mx() - else: - biorbd_return = self.model.InverseDynamics( - q_biorbd, qdot_biorbd, qddot_biorbd, self.biorbd_external_forces_set - ).to_mx() - casadi_fun = Function( - "inverse_dynamics", - [self.q, self.qdot, self.qddot, self.external_forces, self.parameters], - [biorbd_return], - ["q", "qdot", "qddot", "external_forces", "parameters"], - ["tau"], - ) - return casadi_fun - - def contact_forces_from_constrained_forward_dynamics(self) -> Function: - - q_biorbd = GeneralizedCoordinates(self.q) - qdot_biorbd = GeneralizedVelocity(self.qdot) - tau_biorbd = GeneralizedTorque(self.tau) - if self.external_force_set is None: - biorbd_return = self.model.ContactForcesFromForwardDynamicsConstraintsDirect( - q_biorbd, qdot_biorbd, tau_biorbd - ).to_mx() - else: - biorbd_return = self.model.ContactForcesFromForwardDynamicsConstraintsDirect( - q_biorbd, qdot_biorbd, tau_biorbd, self.biorbd_external_forces_set - ).to_mx() - casadi_fun = Function( - "contact_forces_from_constrained_forward_dynamics", + return Function( + "forward_dynamics", [self.q, self.qdot, self.tau, self.external_forces, self.parameters], - [biorbd_return], + [self._dynamic_model(vertcat(self.q, self.qdot, self.tau).T).T], ["q", "qdot", "tau", "external_forces", "parameters"], - ["contact_forces"], - ) - return casadi_fun - - def qdot_from_impact(self) -> Function: - q_biorbd = GeneralizedCoordinates(self.q) - qdot_pre_impact_biorbd = GeneralizedVelocity(self.qdot) - biorbd_return = self.model.ComputeConstraintImpulsesDirect(q_biorbd, qdot_pre_impact_biorbd).to_mx() - casadi_fun = Function( - "qdot_from_impact", - [self.q, self.qdot, self.parameters], - [biorbd_return], - ["q", "qdot", "parameters"], - ["qdot post impact"], - ) - return casadi_fun - - def muscle_activation_dot(self) -> Function: - muscle_states = self.model.stateSet() - for k in range(self.model.nbMuscles()): - muscle_states[k].setActivation(self.activations[k]) - muscle_states[k].setExcitation(self.muscle[k]) - biorbd_return = self.model.activationDot(muscle_states).to_mx() - casadi_fun = Function( - "muscle_activation_dot", - [self.muscle, self.activations, self.parameters], - [biorbd_return], - ["muscle_excitation", "muscle_activation", "parameters"], - ["muscle_activation_dot"], - ) - return casadi_fun - - def muscle_length_jacobian(self) -> Function: - q_biorbd = GeneralizedCoordinates(self.q) - biorbd_return = self.model.musclesLengthJacobian(q_biorbd).to_mx() - casadi_fun = Function( - "muscle_length_jacobian", - [self.q, self.parameters], - [biorbd_return], - ["q", "parameters"], - ["muscle_length_jacobian"], - ) - return casadi_fun - - def muscle_velocity(self) -> Function: - J = self.muscle_length_jacobian()(self.q, self.parameters) - biorbd_return = J @ self.qdot - casadi_fun = Function( - "muscle_velocity", - [self.q, self.qdot, self.parameters], - [biorbd_return], - ["q", "qdot", "parameters"], - ["muscle_velocity"], - ) - return casadi_fun - - def muscle_joint_torque(self) -> Function: - muscles_states = self.model.stateSet() - muscles_activations = self.muscle - for k in range(self.model.nbMuscles()): - muscles_states[k].setActivation(muscles_activations[k]) - q_biorbd = GeneralizedCoordinates(self.q) - qdot_biorbd = GeneralizedVelocity(self.qdot) - biorbd_return = self.model.muscularJointTorque(muscles_states, q_biorbd, qdot_biorbd).to_mx() - casadi_fun = Function( - "muscle_joint_torque", - [self.muscle, self.q, self.qdot, self.parameters], - [biorbd_return], - ["muscle_activation", "q", "qdot", "parameters"], - ["muscle_joint_torque"], - ) - return casadi_fun - - def markers(self) -> list[MX]: - biorbd_return = horzcat(*[m.to_mx() for m in self.model.markers(GeneralizedCoordinates(self.q))]) - casadi_fun = Function( - "markers", - [self.q, self.parameters], - [biorbd_return], - ["q", "parameters"], - ["markers"], - ) - return casadi_fun - - @property - def nb_markers(self) -> int: - return self.model.nbMarkers() - - def marker_index(self, name): - return biorbd.marker_index(self.model, name) - - def marker(self, index: int, reference_segment_index: int = None) -> Function: - q_biorbd = GeneralizedCoordinates(self.q) - marker = self.model.marker(q_biorbd, index) - if reference_segment_index is not None: - global_homogeneous_matrix = self.model.globalJCS(q_biorbd, reference_segment_index) - marker_rotated = global_homogeneous_matrix.transpose().to_mx() @ vertcat(marker.to_mx(), 1) - biorbd_return = marker_rotated[:3] - else: - biorbd_return = marker.to_mx() - casadi_fun = Function( - "marker", - [self.q, self.parameters], - [biorbd_return], - ["q", "parameters"], - ["marker"], - ) - return casadi_fun - - @property - def nb_rigid_contacts(self) -> int: - """ - Returns the number of rigid contacts. - Example: - First contact with axis YZ - Second contact with axis Z - nb_rigid_contacts = 2 - """ - return self.model.nbRigidContacts() + ["qddot"], + ).expand() @property def nb_contacts(self) -> int: - """ - Returns the number of contact index. - Example: - First contact with axis YZ - Second contact with axis Z - nb_contacts = 3 - """ - return self.model.nbContacts() - - def rigid_contact_index(self, contact_index) -> tuple: - """ - Returns the axis index of this specific rigid contact. - Example: - First contact with axis YZ - Second contact with axis Z - rigid_contact_index(0) = (1, 2) - """ - return self.model.rigidContacts()[contact_index].availableAxesIndices() - - def markers_velocities(self, reference_index=None) -> list[MX]: - if reference_index is None: - biorbd_return = [ - m.to_mx() - for m in self.model.markersVelocity( - GeneralizedCoordinates(self.q), - GeneralizedVelocity(self.qdot), - True, - ) - ] - - else: - biorbd_return = [] - homogeneous_matrix_transposed = self.homogeneous_matrices_in_global( - segment_index=reference_index, inverse=True - )( - GeneralizedCoordinates(self.q), - ) - for m in self.model.markersVelocity(GeneralizedCoordinates(self.q), GeneralizedVelocity(self.qdot)): - if m.applyRT(homogeneous_matrix_transposed) is None: - biorbd_return.append(m.to_mx()) - else: - biorbd_return.append(m.applyRT(homogeneous_matrix_transposed).to_mx()) - - casadi_fun = Function( - "markers_velocities", - [self.q, self.qdot, self.parameters], - [horzcat(*biorbd_return)], - ["q", "qdot", "parameters"], - ["markers_velocities"], - ) - return casadi_fun - - def marker_velocity(self, marker_index: int) -> list[MX]: - biorbd_return = self.model.markersVelocity( - GeneralizedCoordinates(self.q), - GeneralizedVelocity(self.qdot), - True, - )[marker_index].to_mx() - casadi_fun = Function( - "marker_velocity", - [self.q, self.qdot, self.parameters], - [biorbd_return], - ["q", "qdot", "parameters"], - ["marker_velocity"], - ) - return casadi_fun - - def markers_accelerations(self, reference_index=None) -> list[MX]: - if reference_index is None: - biorbd_return = [ - m.to_mx() - for m in self.model.markerAcceleration( - GeneralizedCoordinates(self.q), - GeneralizedVelocity(self.qdot), - GeneralizedAcceleration(self.qddot), - True, - ) - ] - - else: - biorbd_return = [] - homogeneous_matrix_transposed = self.homogeneous_matrices_in_global( - segment_index=reference_index, - inverse=True, - )( - GeneralizedCoordinates(self.q), - ) - for m in self.model.markersAcceleration( - GeneralizedCoordinates(self.q), - GeneralizedVelocity(self.qdot), - GeneralizedAcceleration(self.qddot), - ): - if m.applyRT(homogeneous_matrix_transposed) is None: - biorbd_return.append(m.to_mx()) - else: - biorbd_return.append(m.applyRT(homogeneous_matrix_transposed).to_mx()) + return 0 - casadi_fun = Function( - "markers_accelerations", - [self.q, self.qdot, self.qddot, self.parameters], - [horzcat(*biorbd_return)], - ["q", "qdot", "qddot", "parameters"], - ["markers_accelerations"], - ) - return casadi_fun - - def marker_acceleration(self, marker_index: int) -> list[MX]: - biorbd_return = self.model.markerAcceleration( - GeneralizedCoordinates(self.q), - GeneralizedVelocity(self.qdot), - GeneralizedAcceleration(self.qddot), - True, - )[marker_index].to_mx() - casadi_fun = Function( - "marker_acceleration", - [self.q, self.qdot, self.qddot, self.parameters], - [biorbd_return], - ["q", "qdot", "qddot", "parameters"], - ["marker_acceleration"], - ) - return casadi_fun + @property + def nb_soft_contacts(self) -> int: + return 0 - def tau_max(self) -> tuple[MX, MX]: - self.model.closeActuator() - q_biorbd = GeneralizedCoordinates(self.q) - qdot_biorbd = GeneralizedVelocity(self.qdot) - torque_max, torque_min = self.model.torqueMax(q_biorbd, qdot_biorbd) - casadi_fun = Function( - "tau_max", + def reshape_qdot(self, k_stab=1) -> Function: + return Function( + "reshape_qdot", [self.q, self.qdot, self.parameters], - [torque_max.to_mx(), torque_min.to_mx()], + [self.qdot], ["q", "qdot", "parameters"], - ["tau_max", "tau_min"], - ) - return casadi_fun - - def rigid_contact_acceleration(self, contact_index, contact_axis) -> Function: - q_biorbd = GeneralizedCoordinates(self.q) - qdot_biorbd = GeneralizedVelocity(self.qdot) - qddot_biorbd = GeneralizedAcceleration(self.qddot) - biorbd_return = self.model.rigidContactAcceleration( - q_biorbd, qdot_biorbd, qddot_biorbd, contact_index, True - ).to_mx()[contact_axis] - casadi_fun = Function( - "rigid_contact_acceleration", - [self.q, self.qdot, self.qddot, self.parameters], - [biorbd_return], - ["q", "qdot", "qddot", "parameters"], - ["rigid_contact_acceleration"], - ) - return casadi_fun - - def markers_jacobian(self) -> list[MX]: - biorbd_return = [m.to_mx() for m in self.model.markersJacobian(GeneralizedCoordinates(self.q))] - casadi_fun = Function( - "markers_jacobian", - [self.q, self.parameters], - biorbd_return, - ["q", "parameters"], - ["markers_jacobian"], - ) - return casadi_fun - - @property - def marker_names(self) -> tuple[str, ...]: - return tuple([s.to_string() for s in self.model.markerNames()]) + ["Reshaped qdot"], + ).expand() def soft_contact_forces(self) -> Function: - q_biorbd = GeneralizedCoordinates(self.q) - qdot_biorbd = GeneralizedVelocity(self.qdot) - - biorbd_return = MX.zeros(self.nb_soft_contacts * 6, 1) - for i_sc in range(self.nb_soft_contacts): - soft_contact = self.soft_contact(i_sc) - biorbd_return[i_sc * 6 : (i_sc + 1) * 6, :] = ( - biorbd.SoftContactSphere(soft_contact).computeForceAtOrigin(self.model, q_biorbd, qdot_biorbd).to_mx() - ) - - casadi_fun = Function( + return Function( "soft_contact_forces", [self.q, self.qdot, self.parameters], - [biorbd_return], - ["q", "qdot", "parameters"], - ["soft_contact_forces"], - ) - return casadi_fun - - def normalize_state_quaternions(self) -> Function: - - quat_idx = self.get_quaternion_idx() - biorbd_return = MX.zeros(self.nb_q) - biorbd_return[:] = self.q - - # Normalize quaternion, if needed - for j in range(self.nb_quaternions): - quaternion = vertcat( - self.q[quat_idx[j][3]], - self.q[quat_idx[j][0]], - self.q[quat_idx[j][1]], - self.q[quat_idx[j][2]], - ) - quaternion /= norm_fro(quaternion) - biorbd_return[quat_idx[j][0] : quat_idx[j][2] + 1] = quaternion[1:4] - biorbd_return[quat_idx[j][3]] = quaternion[0] - - casadi_fun = Function( - "normalize_state_quaternions", - [self.q], - [biorbd_return], - ["q"], - ["q_normalized"], - ) - return casadi_fun - - def get_quaternion_idx(self) -> list[list[int]]: - n_dof = 0 - quat_idx = [] - quat_number = 0 - for j in range(self.nb_segments): - if self.segments[j].isRotationAQuaternion(): - quat_idx.append([n_dof, n_dof + 1, n_dof + 2, self.nb_dof + quat_number]) - quat_number += 1 - n_dof += self.segments[j].nbDof() - return quat_idx - - def contact_forces(self) -> Function: - force = self.contact_forces_from_constrained_forward_dynamics()( - self.q, self.qdot, self.tau, self.external_forces, self.parameters - ) - casadi_fun = Function( - "contact_forces", - [self.q, self.qdot, self.tau, self.external_forces, self.parameters], - [force], - ["q", "qdot", "tau", "external_forces", "parameters"], - ["contact_forces"], - ) - return casadi_fun - - def passive_joint_torque(self) -> Function: - q_biorbd = GeneralizedCoordinates(self.q) - qdot_biorbd = GeneralizedVelocity(self.qdot) - biorbd_return = self.model.passiveJointTorque(q_biorbd, qdot_biorbd).to_mx() - casadi_fun = Function( - "passive_joint_torque", - [self.q, self.qdot, self.parameters], - [biorbd_return], - ["q", "qdot", "parameters"], - ["passive_joint_torque"], - ) - return casadi_fun - - def ligament_joint_torque(self) -> Function: - q_biorbd = GeneralizedCoordinates(self.q) - qdot_biorbd = GeneralizedVelocity(self.qdot) - biorbd_return = self.model.ligamentsJointTorque(q_biorbd, qdot_biorbd).to_mx() - casadi_fun = Function( - "ligament_joint_torque", - [self.q, self.qdot, self.parameters], - [biorbd_return], + [MX(0)], ["q", "qdot", "parameters"], - ["ligament_joint_torque"], - ) - return casadi_fun - - def ranges_from_model(self, variable: str): - ranges = [] - for segment in self.segments: - if "_joints" in variable: - if segment.parent().to_string().lower() != "root": - variable = variable.replace("_joints", "") - ranges += self._add_range(variable, segment) - elif "_roots" in variable: - if segment.parent().to_string().lower() == "root": - variable = variable.replace("_roots", "") - ranges += self._add_range(variable, segment) - else: - ranges += self._add_range(variable, segment) - - return ranges - - @staticmethod - def _add_range(variable: str, segment: biorbd.Segment) -> list[biorbd.Range]: - """ - Get the range of a variable for a given segment - - Parameters - ---------- - variable: str - The variable to get the range for such as: - "q", "qdot", "qddot", "q_joint", "qdot_joint", "qddot_joint", "q_root", "qdot_root", "qddot_root" - segment: biorbd.Segment - The segment to get the range from - - Returns - ------- - list[biorbd.Range] - range min and max for the given variable for a given segment - """ - ranges_map = { - "q": [q_range for q_range in segment.QRanges()], - "qdot": [qdot_range for qdot_range in segment.QdotRanges()], - "qddot": [qddot_range for qddot_range in segment.QddotRanges()], - } - - segment_variable_range = ranges_map.get(variable, None) - if segment_variable_range is None: - RuntimeError("Wrong variable name") - - return segment_variable_range - - def _var_mapping( - self, - key: str, - range_for_mapping: int | list | tuple | range, - mapping: BiMapping = None, - ) -> dict: - return _var_mapping(key, range_for_mapping, mapping) - - def bounds_from_ranges(self, variables: str | list[str], mapping: BiMapping | BiMappingList = None) -> Bounds: - return bounds_from_ranges(self, variables, mapping) - - def lagrangian(self) -> Function: - q_biorbd = GeneralizedCoordinates(self.q) - qdot_biorbd = GeneralizedVelocity(self.qdot) - biorbd_return = self.model.Lagrangian(q_biorbd, qdot_biorbd).to_mx() - casadi_fun = Function( - "lagrangian", - [self.q, self.qdot], - [biorbd_return], - ["q", "qdot"], - ["lagrangian"], - ) - return casadi_fun - - def partitioned_forward_dynamics(self): - raise NotImplementedError("partitioned_forward_dynamics is not implemented for BiorbdModel") - - @staticmethod - def animate( - ocp, - solution, - show_now: bool = True, - show_tracked_markers: bool = False, - viewer: str = "pyorerun", - n_frames: int = 0, - **kwargs, - ): - if viewer == "bioviz": - from .viewer_bioviz import animate_with_bioviz_for_loop - - return animate_with_bioviz_for_loop(ocp, solution, show_now, show_tracked_markers, n_frames, **kwargs) - if viewer == "pyorerun": - from .viewer_pyorerun import animate_with_pyorerun - - return animate_with_pyorerun(ocp, solution, show_now, show_tracked_markers, **kwargs) + ["Soft contact forces"], + ).expand() From daefed02cafdc526951dabce27e2aabfa7a2da70 Mon Sep 17 00:00:00 2001 From: Pariterre Date: Wed, 18 Dec 2024 07:37:37 -0500 Subject: [PATCH 05/19] Trying something --- .gitignore | 2 +- .../custom_non_casadi_dynamics.py | 44 +++++++++++++++---- bioptim/models/torch/torch_model.py | 2 +- 3 files changed, 38 insertions(+), 10 deletions(-) diff --git a/.gitignore b/.gitignore index 12a9d2595..b5f3cce77 100644 --- a/.gitignore +++ b/.gitignore @@ -122,7 +122,7 @@ _l4c_generated/ *.png -bioptim/examples/getting_started/_l4c_generated/ +bioptim/*/_l4c_generated/ sandbox/ # Ignore all npy files in tests folder except those in tests/shard6 diff --git a/bioptim/examples/getting_started/custom_non_casadi_dynamics.py b/bioptim/examples/getting_started/custom_non_casadi_dynamics.py index a32513ede..efe586fe5 100644 --- a/bioptim/examples/getting_started/custom_non_casadi_dynamics.py +++ b/bioptim/examples/getting_started/custom_non_casadi_dynamics.py @@ -71,7 +71,7 @@ def train_me(self, training_data: list[torch.Tensor], validation_data: list[torc self._optimizer, mode="min", factor=0.1, patience=20, min_lr=1e-8 ) - max_epochs = 10 + max_epochs = 500 for _ in range(max_epochs): self._perform_epoch_training(targets=training_data) validation_loss = self._perform_epoch_training(targets=validation_data, only_compute=True) @@ -130,6 +130,13 @@ def prepare_ocp( dynamics = Dynamics(DynamicsFcn.TORQUE_DRIVEN) torch_model = TorchModel(torch_model=model) + q = np.array([0, 3.14]) + qdot = np.array([0, 0]) + tau = np.array([0, 0]) + qddot = torch_model.forward_dynamics()(q, qdot, tau, np.array([]), np.array([])) + biorbd_model = biorbd.Model("models/pendulum.bioMod") + qddot2 = biorbd_model.ForwardDynamics(q, qdot, tau).to_array() + # Path bounds x_bounds = BoundsList() x_bounds["q"] = [-3.14 * 1.5] * torch_model.nb_q, [3.14 * 1.5] * torch_model.nb_q @@ -150,14 +157,35 @@ def prepare_ocp( final_time, x_bounds=x_bounds, u_bounds=u_bounds, - use_sx=True, + use_sx=False, ) def generate_dataset(biorbd_model: biorbd.Model, data_point_count: int) -> list[torch.Tensor]: - q = torch.rand(data_point_count, biorbd_model.nbQ()) - qdot = torch.rand(data_point_count, biorbd_model.nbQdot()) - tau = torch.rand(data_point_count, biorbd_model.nbGeneralizedTorque()) + q_ranges = np.array( + [[[q_range.min(), q_range.max()] for q_range in segment.QRanges()] for segment in biorbd_model.segments()] + ).squeeze() + qdot_ranges = np.array( + [ + [[qdot_range.min(), qdot_range.max()] for qdot_range in segment.QdotRanges()] + for segment in biorbd_model.segments() + ] + ).squeeze() + tau_ranges = np.array([-100, 100] * biorbd_model.nbGeneralizedTorque()).reshape(-1, 2) + + q = torch.rand(data_point_count, biorbd_model.nbQ()) * (q_ranges[:, 1] - q_ranges[:, 0]) + q_ranges[:, 0] + qdot = ( + torch.rand(data_point_count, biorbd_model.nbQdot()) * (qdot_ranges[:, 1] - qdot_ranges[:, 0]) + + qdot_ranges[:, 0] + ) + tau = ( + torch.rand(data_point_count, biorbd_model.nbGeneralizedTorque()) * (tau_ranges[:, 1] - tau_ranges[:, 0]) + + tau_ranges[:, 0] + ) + + q = q.to(torch.float) + qdot = qdot.to(torch.float) + tau = tau.to(torch.float) qddot = torch.zeros(data_point_count, biorbd_model.nbQddot()) for i in range(data_point_count): @@ -171,10 +199,10 @@ def generate_dataset(biorbd_model: biorbd.Model, data_point_count: int) -> list[ def main(): # --- Prepare a predictive model --- # biorbd_model = biorbd.Model("models/pendulum.bioMod") - training_data = generate_dataset(biorbd_model, data_point_count=1000) - validation_data = generate_dataset(biorbd_model, data_point_count=100) + training_data = generate_dataset(biorbd_model, data_point_count=30000) + validation_data = generate_dataset(biorbd_model, data_point_count=3000) - model = NeuralNetworkModel(layer_node_count=(6, 10, 10, 2), dropout_probability=0.2, use_batch_norm=True) + model = NeuralNetworkModel(layer_node_count=(6, 128, 128, 128, 2), dropout_probability=0.2, use_batch_norm=True) model.train_me(training_data, validation_data) ocp = prepare_ocp(model=model, final_time=1, n_shooting=40) diff --git a/bioptim/models/torch/torch_model.py b/bioptim/models/torch/torch_model.py index abdfcf28b..585170b88 100644 --- a/bioptim/models/torch/torch_model.py +++ b/bioptim/models/torch/torch_model.py @@ -72,7 +72,7 @@ def forward_dynamics(self, with_contact: bool = False) -> Function: [self._dynamic_model(vertcat(self.q, self.qdot, self.tau).T).T], ["q", "qdot", "tau", "external_forces", "parameters"], ["qddot"], - ).expand() + ) @property def nb_contacts(self) -> int: From 9766eaf47ae959e1d54de5025d653c35fe880e8f Mon Sep 17 00:00:00 2001 From: Pariterre Date: Thu, 9 Jan 2025 14:59:07 -0500 Subject: [PATCH 06/19] Finished implementing for L4Casadi on Windows. Still can't use with Ipopt --- .gitignore | 7 +- .../custom_non_casadi_dynamics.py | 90 +++++++++++++++++-- 2 files changed, 82 insertions(+), 15 deletions(-) diff --git a/.gitignore b/.gitignore index b5f3cce77..3b0a16dd5 100644 --- a/.gitignore +++ b/.gitignore @@ -109,7 +109,7 @@ c_generated_code *.gv.pdf # Bioptim files -bioptim/sandbox/ +sandbox/ *.pkl *.mtx *.casadi @@ -120,11 +120,6 @@ _l4c_generated/ # Mac dev *.DS_store -*.png - -bioptim/*/_l4c_generated/ -sandbox/ - # Ignore all npy files in tests folder except those in tests/shard6 /tests/*.npy !/tests/shard6/v_bounds_max.npy diff --git a/bioptim/examples/getting_started/custom_non_casadi_dynamics.py b/bioptim/examples/getting_started/custom_non_casadi_dynamics.py index efe586fe5..0906879a2 100644 --- a/bioptim/examples/getting_started/custom_non_casadi_dynamics.py +++ b/bioptim/examples/getting_started/custom_non_casadi_dynamics.py @@ -6,6 +6,9 @@ based model. This is useful when the dynamics are computed using a different library (e.g. TensorFlow, PyTorch, etc.) """ +import os +from typing import Self + import biorbd from bioptim import OptimalControlProgram, DynamicsFcn, BoundsList, Dynamics from bioptim.models.torch.torch_model import TorchModel @@ -13,6 +16,24 @@ import torch +class EarlyStopper: + def __init__(self, patience=1, min_delta=0): + self.patience = patience + self.min_delta = min_delta + self.counter = 0 + self.min_validation_loss = float("inf") + + def early_stop(self, validation_loss): + if validation_loss < self.min_validation_loss: + self.min_validation_loss = validation_loss + self.counter = 0 + elif validation_loss > (self.min_validation_loss + self.min_delta): + self.counter += 1 + if self.counter >= self.patience: + return True + return False + + class NeuralNetworkModel(torch.nn.Module): def __init__( self, @@ -33,7 +54,8 @@ def __init__( torch.nn.Linear(first_and_hidden_layers_node_count[i], first_and_hidden_layers_node_count[i + 1]) ) if use_batch_norm: - torch.nn.BatchNorm1d(first_and_hidden_layers_node_count[i + 1]) + raise NotImplementedError("Batch normalization is not yet implemented") + layers.append(torch.nn.BatchNorm1d(first_and_hidden_layers_node_count[i + 1])) layers.append(activations) layers.append(torch.nn.Dropout(dropout_probability)) layers.append(torch.nn.Linear(first_and_hidden_layers_node_count[-1], layer_node_count[-1])) @@ -65,19 +87,63 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: def get_torch_device() -> torch.device: return torch.device("cuda" if torch.cuda.is_available() else "cpu") - def train_me(self, training_data: list[torch.Tensor], validation_data: list[torch.Tensor]): + def train_me( + self, training_data: list[torch.Tensor], validation_data: list[torch.Tensor], max_epochs: int = 5 + ) -> None: # More details about scheduler in documentation scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( self._optimizer, mode="min", factor=0.1, patience=20, min_lr=1e-8 ) - max_epochs = 500 - for _ in range(max_epochs): + early_stopper = EarlyStopper(patience=20, min_delta=1e-5) + for i in range(max_epochs): self._perform_epoch_training(targets=training_data) validation_loss = self._perform_epoch_training(targets=validation_data, only_compute=True) - print(f"Validation loss: {validation_loss}") scheduler.step(validation_loss) # Adjust/reduce learning rate + # Check if the training should stop + print(f"Validation loss: {validation_loss} (epoch: {i})") + if early_stopper.early_stop(validation_loss): + print("Early stopping") + break + + def save_me(self, path: str) -> None: + layer_node_count = tuple( + [model.in_features for model in self._forward_model if isinstance(model, torch.nn.Linear)] + + [self._forward_model[-1].out_features] + ) + + dropout_probability = tuple([model.p for model in self._forward_model if isinstance(model, torch.nn.Dropout)]) + if len(dropout_probability) == 0: + dropout_probability = 0 + elif len(dropout_probability) > 1: + # make sure that the dropout probability is the same for all layers + if not all(prob == dropout_probability[0] for prob in dropout_probability): + raise ValueError("Different dropout probabilities for different layers") + dropout_probability = dropout_probability[0] + + use_batch_norm = any(isinstance(model, torch.nn.BatchNorm1d) for model in self._forward_model) + + dico = { + "layer_node_count": layer_node_count, + "dropout_probability": dropout_probability, + "use_batch_norm": use_batch_norm, + "state_dict": self.state_dict(), + } + torch.save(dico, path) + + @classmethod + def load_me(cls, path: str) -> Self: + data = torch.load(path, weights_only=True) + inputs = { + "layer_node_count": data["layer_node_count"], + "dropout_probability": data["dropout_probability"], + "use_batch_norm": data["use_batch_norm"], + } + model = NeuralNetworkModel(**inputs) + model.load_state_dict(data["state_dict"]) + return model + def _perform_epoch_training( self, targets: list[torch.Tensor], @@ -130,12 +196,13 @@ def prepare_ocp( dynamics = Dynamics(DynamicsFcn.TORQUE_DRIVEN) torch_model = TorchModel(torch_model=model) - q = np.array([0, 3.14]) + q = np.array([0, 0]) qdot = np.array([0, 0]) tau = np.array([0, 0]) - qddot = torch_model.forward_dynamics()(q, qdot, tau, np.array([]), np.array([])) + qddot = torch_model.forward_dynamics()(q, qdot, tau, [], []) biorbd_model = biorbd.Model("models/pendulum.bioMod") qddot2 = biorbd_model.ForwardDynamics(q, qdot, tau).to_array() + print(qddot - qddot2) # Path bounds x_bounds = BoundsList() @@ -198,12 +265,17 @@ def generate_dataset(biorbd_model: biorbd.Model, data_point_count: int) -> list[ def main(): # --- Prepare a predictive model --- # + force_new_training = False biorbd_model = biorbd.Model("models/pendulum.bioMod") training_data = generate_dataset(biorbd_model, data_point_count=30000) validation_data = generate_dataset(biorbd_model, data_point_count=3000) - model = NeuralNetworkModel(layer_node_count=(6, 128, 128, 128, 2), dropout_probability=0.2, use_batch_norm=True) - model.train_me(training_data, validation_data) + if force_new_training or not os.path.isfile("models/trained_pendulum_model.pt"): + model = NeuralNetworkModel(layer_node_count=(6, 512, 512, 2), dropout_probability=0.2, use_batch_norm=False) + model.train_me(training_data, validation_data, max_epochs=300) + model.save_me("models/trained_pendulum_model.pt") + else: + model = NeuralNetworkModel.load_me("models/trained_pendulum_model.pt") ocp = prepare_ocp(model=model, final_time=1, n_shooting=40) From 70edeef59f7f1582862ddc5633f9f17b2e4dfd9b Mon Sep 17 00:00:00 2001 From: Pariterre Date: Thu, 9 Jan 2025 16:47:51 -0500 Subject: [PATCH 07/19] Added dummy non working example --- bioptim/examples/getting_started/race_car.py | 65 ++++++++++++++++++++ 1 file changed, 65 insertions(+) create mode 100644 bioptim/examples/getting_started/race_car.py diff --git a/bioptim/examples/getting_started/race_car.py b/bioptim/examples/getting_started/race_car.py new file mode 100644 index 000000000..4ef6e9995 --- /dev/null +++ b/bioptim/examples/getting_started/race_car.py @@ -0,0 +1,65 @@ +from casadi import * +import l4casadi as l4c +import torch + + +class NeuralNetworkModel(torch.nn.Module): + def __init__(self, layer_node_count: tuple[int]): + super(NeuralNetworkModel, self).__init__() + + # Initialize the layers of the neural network + self._size_in = layer_node_count[0] + self._size_out = layer_node_count[-1] + first_and_hidden_layers_node_count = layer_node_count[:-1] + layers = torch.nn.ModuleList() + for i in range(len(first_and_hidden_layers_node_count) - 1): + layers.append( + torch.nn.Linear(first_and_hidden_layers_node_count[i], first_and_hidden_layers_node_count[i + 1]) + ) + layers.append(torch.nn.Linear(first_and_hidden_layers_node_count[-1], layer_node_count[-1])) + + self._forward_model = torch.nn.Sequential(*layers) + self._forward_model.to("cpu") + + # Put the model in evaluation mode + self.eval() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + output = torch.Tensor(x.shape[0], self._forward_model[-1].out_features) + for i, data in enumerate(x): + output[i, :] = self._forward_model(data) + return output.to("cpu") + + +def main(): + opti = Opti() # Optimization problem + nx = nu = 1 + hidden_layers = (10,) + N = 100 # number of control intervals + + # ---- decision variables --------- + X = opti.variable(nx, N + 1) # state trajectory + U = opti.variable(nu, N) # control trajectory (throttle) + + # ---- dynamic constraints -------- + torch_model = NeuralNetworkModel(layer_node_count=(nx + nu, *hidden_layers, nx)) + dynamic_model = l4c.L4CasADi(torch_model, device="cpu") + x_sym = MX.sym("x", nx, 1) + u_sym = MX.sym("u", nu, 1) + f = Function("qddot", [x_sym, u_sym], [dynamic_model(vertcat(x_sym, u_sym).T).T]) + + for k in range(N): # loop over control intervals + # Runge-Kutta 1 integration + k1 = f(X[:, k], U[:, k]) + opti.subject_to(X[:, k + 1] == k1) # close the gaps + + # ---- boundary conditions -------- + opti.subject_to(X[0, 0] == 1) # PROBLEM LIES HERE: PUTTING ANY VALUE BUT 0 WILL MAKE IPOPT FAILS + + # ---- solve NLP ------ + opti.solver("ipopt") # set numerical backend + opti.solve() # actual solve + + +if __name__ == "__main__": + main() From 31635a8614a2226d00ceee7fb013e98b5c444cb9 Mon Sep 17 00:00:00 2001 From: Pariterre Date: Fri, 10 Jan 2025 09:22:01 -0500 Subject: [PATCH 08/19] Made the most minimal version of the non-working l4casadi example --- bioptim/examples/getting_started/race_car.py | 65 ------------------- .../temporary_non_working_example.py | 40 ++++++++++++ 2 files changed, 40 insertions(+), 65 deletions(-) delete mode 100644 bioptim/examples/getting_started/race_car.py create mode 100644 bioptim/examples/getting_started/temporary_non_working_example.py diff --git a/bioptim/examples/getting_started/race_car.py b/bioptim/examples/getting_started/race_car.py deleted file mode 100644 index 4ef6e9995..000000000 --- a/bioptim/examples/getting_started/race_car.py +++ /dev/null @@ -1,65 +0,0 @@ -from casadi import * -import l4casadi as l4c -import torch - - -class NeuralNetworkModel(torch.nn.Module): - def __init__(self, layer_node_count: tuple[int]): - super(NeuralNetworkModel, self).__init__() - - # Initialize the layers of the neural network - self._size_in = layer_node_count[0] - self._size_out = layer_node_count[-1] - first_and_hidden_layers_node_count = layer_node_count[:-1] - layers = torch.nn.ModuleList() - for i in range(len(first_and_hidden_layers_node_count) - 1): - layers.append( - torch.nn.Linear(first_and_hidden_layers_node_count[i], first_and_hidden_layers_node_count[i + 1]) - ) - layers.append(torch.nn.Linear(first_and_hidden_layers_node_count[-1], layer_node_count[-1])) - - self._forward_model = torch.nn.Sequential(*layers) - self._forward_model.to("cpu") - - # Put the model in evaluation mode - self.eval() - - def forward(self, x: torch.Tensor) -> torch.Tensor: - output = torch.Tensor(x.shape[0], self._forward_model[-1].out_features) - for i, data in enumerate(x): - output[i, :] = self._forward_model(data) - return output.to("cpu") - - -def main(): - opti = Opti() # Optimization problem - nx = nu = 1 - hidden_layers = (10,) - N = 100 # number of control intervals - - # ---- decision variables --------- - X = opti.variable(nx, N + 1) # state trajectory - U = opti.variable(nu, N) # control trajectory (throttle) - - # ---- dynamic constraints -------- - torch_model = NeuralNetworkModel(layer_node_count=(nx + nu, *hidden_layers, nx)) - dynamic_model = l4c.L4CasADi(torch_model, device="cpu") - x_sym = MX.sym("x", nx, 1) - u_sym = MX.sym("u", nu, 1) - f = Function("qddot", [x_sym, u_sym], [dynamic_model(vertcat(x_sym, u_sym).T).T]) - - for k in range(N): # loop over control intervals - # Runge-Kutta 1 integration - k1 = f(X[:, k], U[:, k]) - opti.subject_to(X[:, k + 1] == k1) # close the gaps - - # ---- boundary conditions -------- - opti.subject_to(X[0, 0] == 1) # PROBLEM LIES HERE: PUTTING ANY VALUE BUT 0 WILL MAKE IPOPT FAILS - - # ---- solve NLP ------ - opti.solver("ipopt") # set numerical backend - opti.solve() # actual solve - - -if __name__ == "__main__": - main() diff --git a/bioptim/examples/getting_started/temporary_non_working_example.py b/bioptim/examples/getting_started/temporary_non_working_example.py new file mode 100644 index 000000000..762ac8735 --- /dev/null +++ b/bioptim/examples/getting_started/temporary_non_working_example.py @@ -0,0 +1,40 @@ +from casadi import Opti, MX, Function +import l4casadi as l4c +import torch + + +class NeuralNetworkModel(torch.nn.Module): + def __init__(self, layer_node_count: tuple[int]): + super(NeuralNetworkModel, self).__init__() + layers = torch.nn.ModuleList() + layers.append(torch.nn.Linear(layer_node_count[0], layer_node_count[-1])) + self._forward_model = torch.nn.Sequential(*layers) + self.eval() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self._forward_model(x) + + +def main(): + opti = Opti() + nx = nu = 1 + torch_model = NeuralNetworkModel(layer_node_count=(nu, nx)) + + # ---- decision variables --------- + x = opti.variable(nx, 1) # state + u = opti.variable(nu, 1) # control + + # ---- dynamic constraints -------- + x_sym = MX.sym("x", nx, 1) + u_sym = MX.sym("u", nu, 1) + forward_model = l4c.L4CasADi(torch_model, device="cpu") + f = Function("xdot", [x_sym, u_sym], [x_sym - forward_model(u_sym)]) + opti.subject_to(f(x, u) == 0) # Adding this line yields the error : jac_adj_i0_adj_o0 is not provided by L4CasADi. + + # ---- solve NLP ------ + opti.solver("ipopt") + opti.solve() + + +if __name__ == "__main__": + main() From 320696863cea92316a7c4e5a9a09f522e9a3bac8 Mon Sep 17 00:00:00 2001 From: Pariterre Date: Mon, 13 Jan 2025 15:55:15 -0500 Subject: [PATCH 09/19] Finalized the torch example and added a test --- .github/workflows/run_tests_linux.yml | 8 ++ .gitmodules | 3 + bioptim/examples/__main__.py | 8 ++ .../pytorch_ocp.py} | 60 +++++++-------- .../temporary_non_working_example.py | 40 ---------- bioptim/interfaces/interface_utils.py | 2 + bioptim/models/torch/torch_model.py | 74 +++++++++++++------ external/l4casadi | 1 + tests/shard3/test_global_getting_started.py | 14 ++++ 9 files changed, 114 insertions(+), 96 deletions(-) rename bioptim/examples/{getting_started/custom_non_casadi_dynamics.py => deep_neural_network/pytorch_ocp.py} (86%) delete mode 100644 bioptim/examples/getting_started/temporary_non_working_example.py create mode 160000 external/l4casadi diff --git a/.github/workflows/run_tests_linux.yml b/.github/workflows/run_tests_linux.yml index 0f7a02e50..dfc7f273d 100644 --- a/.github/workflows/run_tests_linux.yml +++ b/.github/workflows/run_tests_linux.yml @@ -51,6 +51,14 @@ jobs: python -c "import bioptim" if: matrix.shard == 1 + - name: Install pytorch on Linux + run: | + pip install torch>=2.0 + cd external/l4casadi + pip install . --no-build-isolation + cd ../.. + if: matrix.shard == 3 + - name: Run tests with code coverage run: pytest -v --color=yes --cov-report term-missing --cov=bioptim tests/shard${{ matrix.shard }} --mpl-baseline-path=bioptim/tests/plot_reference_images if: matrix.os == 'ubuntu-latest' diff --git a/.gitmodules b/.gitmodules index 3df878b8e..633e50b64 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,6 @@ [submodule "external/acados"] path = external/acados url = https://github.com/acados/acados.git +[submodule "external/l4casadi"] + path = external/l4casadi + url = https://github.com/pariterre/l4casadi diff --git a/bioptim/examples/__main__.py b/bioptim/examples/__main__.py index 523141908..19b1b2a82 100644 --- a/bioptim/examples/__main__.py +++ b/bioptim/examples/__main__.py @@ -135,6 +135,14 @@ ] ), ), + ( + "deep_neural_network", + OrderedDict( + [ + ("pytorch ocp", "pytorch_ocp.py"), + ] + ), + ), ] ) diff --git a/bioptim/examples/getting_started/custom_non_casadi_dynamics.py b/bioptim/examples/deep_neural_network/pytorch_ocp.py similarity index 86% rename from bioptim/examples/getting_started/custom_non_casadi_dynamics.py rename to bioptim/examples/deep_neural_network/pytorch_ocp.py index 0906879a2..e4a24cb75 100644 --- a/bioptim/examples/getting_started/custom_non_casadi_dynamics.py +++ b/bioptim/examples/deep_neural_network/pytorch_ocp.py @@ -1,16 +1,21 @@ """ -TODO: Explain what is this example about -TODO: All the documentation - -This example is similar to the getting_started/pendulum.py example, but the dynamics are computed using a non-casadi -based model. This is useful when the dynamics are computed using a different library (e.g. TensorFlow, PyTorch, etc.) +This example is similar to the getting_started/pendulum.py example, but the dynamics are computed using a deep neural +network driven by PyTorch. For this example to work, we create a neural network model that trains against the biorbd +library to predict the acceleration of the pendulum. Obviously, one can replace the prediction model with any other +pytorch model. The class TorchModel is thereafter used to wrap the model using L4casadi to be used in the bioptim framework, +if one needs a more flexible way to define the dynamics, objective functions or constraints, they can start from that +class and modify it to their needs. + +Extra information: +All information regarding the installation, limitations, credits and known problems can be found in the header of the +TorchModel class. """ import os from typing import Self import biorbd -from bioptim import OptimalControlProgram, DynamicsFcn, BoundsList, Dynamics +from bioptim import OptimalControlProgram, DynamicsFcn, BoundsList, Dynamics, Objective, ObjectiveFcn, Solver from bioptim.models.torch.torch_model import TorchModel import numpy as np import torch @@ -35,12 +40,7 @@ def early_stop(self, validation_loss): class NeuralNetworkModel(torch.nn.Module): - def __init__( - self, - layer_node_count: tuple[int], - dropout_probability: float, - use_batch_norm: bool, - ): + def __init__(self, layer_node_count: tuple[int], dropout_probability: float): super(NeuralNetworkModel, self).__init__() activations = torch.nn.GELU() @@ -53,9 +53,6 @@ def __init__( layers.append( torch.nn.Linear(first_and_hidden_layers_node_count[i], first_and_hidden_layers_node_count[i + 1]) ) - if use_batch_norm: - raise NotImplementedError("Batch normalization is not yet implemented") - layers.append(torch.nn.BatchNorm1d(first_and_hidden_layers_node_count[i + 1])) layers.append(activations) layers.append(torch.nn.Dropout(dropout_probability)) layers.append(torch.nn.Linear(first_and_hidden_layers_node_count[-1], layer_node_count[-1])) @@ -88,7 +85,7 @@ def get_torch_device() -> torch.device: return torch.device("cuda" if torch.cuda.is_available() else "cpu") def train_me( - self, training_data: list[torch.Tensor], validation_data: list[torch.Tensor], max_epochs: int = 5 + self, training_data: list[torch.Tensor], validation_data: list[torch.Tensor], max_epochs: int = 100 ) -> None: # More details about scheduler in documentation scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( @@ -121,15 +118,16 @@ def save_me(self, path: str) -> None: if not all(prob == dropout_probability[0] for prob in dropout_probability): raise ValueError("Different dropout probabilities for different layers") dropout_probability = dropout_probability[0] - - use_batch_norm = any(isinstance(model, torch.nn.BatchNorm1d) for model in self._forward_model) + else: + dropout_probability = dropout_probability[0] dico = { "layer_node_count": layer_node_count, "dropout_probability": dropout_probability, - "use_batch_norm": use_batch_norm, "state_dict": self.state_dict(), } + if not os.path.isdir(os.path.dirname(path)): + os.makedirs(os.path.dirname(path)) torch.save(dico, path) @classmethod @@ -138,7 +136,6 @@ def load_me(cls, path: str) -> Self: inputs = { "layer_node_count": data["layer_node_count"], "dropout_probability": data["dropout_probability"], - "use_batch_norm": data["use_batch_norm"], } model = NeuralNetworkModel(**inputs) model.load_state_dict(data["state_dict"]) @@ -196,14 +193,6 @@ def prepare_ocp( dynamics = Dynamics(DynamicsFcn.TORQUE_DRIVEN) torch_model = TorchModel(torch_model=model) - q = np.array([0, 0]) - qdot = np.array([0, 0]) - tau = np.array([0, 0]) - qddot = torch_model.forward_dynamics()(q, qdot, tau, [], []) - biorbd_model = biorbd.Model("models/pendulum.bioMod") - qddot2 = biorbd_model.ForwardDynamics(q, qdot, tau).to_array() - print(qddot - qddot2) - # Path bounds x_bounds = BoundsList() x_bounds["q"] = [-3.14 * 1.5] * torch_model.nb_q, [3.14 * 1.5] * torch_model.nb_q @@ -217,6 +206,9 @@ def prepare_ocp( u_bounds["tau"] = [-100] * torch_model.nb_tau, [100] * torch_model.nb_tau u_bounds["tau"][1, :] = 0 # ...but remove the capability to actively rotate + # Add objective functions + objective_functions = Objective(ObjectiveFcn.Lagrange.MINIMIZE_CONTROL, key="tau") + return OptimalControlProgram( torch_model, dynamics, @@ -224,6 +216,7 @@ def prepare_ocp( final_time, x_bounds=x_bounds, u_bounds=u_bounds, + objective_functions=objective_functions, use_sx=False, ) @@ -239,6 +232,7 @@ def generate_dataset(biorbd_model: biorbd.Model, data_point_count: int) -> list[ ] ).squeeze() tau_ranges = np.array([-100, 100] * biorbd_model.nbGeneralizedTorque()).reshape(-1, 2) + tau_ranges[1, :] = 0 q = torch.rand(data_point_count, biorbd_model.nbQ()) * (q_ranges[:, 1] - q_ranges[:, 0]) + q_ranges[:, 0] qdot = ( @@ -266,12 +260,12 @@ def generate_dataset(biorbd_model: biorbd.Model, data_point_count: int) -> list[ def main(): # --- Prepare a predictive model --- # force_new_training = False - biorbd_model = biorbd.Model("models/pendulum.bioMod") + biorbd_model = biorbd.Model("../getting_started/models/pendulum.bioMod") training_data = generate_dataset(biorbd_model, data_point_count=30000) validation_data = generate_dataset(biorbd_model, data_point_count=3000) if force_new_training or not os.path.isfile("models/trained_pendulum_model.pt"): - model = NeuralNetworkModel(layer_node_count=(6, 512, 512, 2), dropout_probability=0.2, use_batch_norm=False) + model = NeuralNetworkModel(layer_node_count=(6, 8, 2), dropout_probability=0.2) model.train_me(training_data, validation_data, max_epochs=300) model.save_me("models/trained_pendulum_model.pt") else: @@ -280,11 +274,7 @@ def main(): ocp = prepare_ocp(model=model, final_time=1, n_shooting=40) # --- Solve the ocp --- # - sol = ocp.solve() - - # --- Show the results graph --- # - # sol.print_cost() - sol.graphs(show_bounds=True) + ocp.solve(Solver.IPOPT(show_online_optim=True)) if __name__ == "__main__": diff --git a/bioptim/examples/getting_started/temporary_non_working_example.py b/bioptim/examples/getting_started/temporary_non_working_example.py deleted file mode 100644 index 762ac8735..000000000 --- a/bioptim/examples/getting_started/temporary_non_working_example.py +++ /dev/null @@ -1,40 +0,0 @@ -from casadi import Opti, MX, Function -import l4casadi as l4c -import torch - - -class NeuralNetworkModel(torch.nn.Module): - def __init__(self, layer_node_count: tuple[int]): - super(NeuralNetworkModel, self).__init__() - layers = torch.nn.ModuleList() - layers.append(torch.nn.Linear(layer_node_count[0], layer_node_count[-1])) - self._forward_model = torch.nn.Sequential(*layers) - self.eval() - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self._forward_model(x) - - -def main(): - opti = Opti() - nx = nu = 1 - torch_model = NeuralNetworkModel(layer_node_count=(nu, nx)) - - # ---- decision variables --------- - x = opti.variable(nx, 1) # state - u = opti.variable(nu, 1) # control - - # ---- dynamic constraints -------- - x_sym = MX.sym("x", nx, 1) - u_sym = MX.sym("u", nu, 1) - forward_model = l4c.L4CasADi(torch_model, device="cpu") - f = Function("xdot", [x_sym, u_sym], [x_sym - forward_model(u_sym)]) - opti.subject_to(f(x, u) == 0) # Adding this line yields the error : jac_adj_i0_adj_o0 is not provided by L4CasADi. - - # ---- solve NLP ------ - opti.solver("ipopt") - opti.solve() - - -if __name__ == "__main__": - main() diff --git a/bioptim/interfaces/interface_utils.py b/bioptim/interfaces/interface_utils.py index 6d91d416b..b8cf51512 100644 --- a/bioptim/interfaces/interface_utils.py +++ b/bioptim/interfaces/interface_utils.py @@ -1,3 +1,4 @@ +import warnings from time import perf_counter from casadi import Importer, Function @@ -214,6 +215,7 @@ def _shake_penalties_tree(ocp, penalties_cx: CX, v: CX, v_bounds: DoubleNpArrayT penalty = penalty.expand() except RuntimeError: # This happens mostly when, for instance, there is a Newton decent in the penalty + warnings.warn("Could not expand during the shake_tree.") pass return penalty(vertcat(*dt, v[len(dt) :])) diff --git a/bioptim/models/torch/torch_model.py b/bioptim/models/torch/torch_model.py index 585170b88..270dbeeea 100644 --- a/bioptim/models/torch/torch_model.py +++ b/bioptim/models/torch/torch_model.py @@ -1,33 +1,62 @@ -from typing import Callable - -from casadi import SX, MX, vertcat, horzcat, norm_fro, Function +""" +This wrapper can be used to wrap a PyTorch model and use it in bioptim. This is a much incomplete class though as compared +to the BiorbdModel class. The main reason is that as opposed to Biorbd, the dynamics produced by a PyTorch model can be +of any nature. This means this wrapper be more viewed as an example of how to wrap a PyTorch model in bioptim than an actual +generic wrapper. + +This wrapper is based on the l4casadi library (https://github.com/Tim-Salzmann/l4casadi) which is a bridge between CasADi +and PyTorch. + +INSTALLATION: +Note these instructions may be outdated. Please refer to the l4casadi documentation for the most up-to-date instructions. + +First, make sure pytorch is installed by running the following command: + pip install torch>=2.0 +Please note that some depencecies are required. At the time of writing, the following packages were required: + pip install setuptools>=68.1 scikit-build>=0.17 cmake>=3.27 ninja>=1.11 +Then, install l4casadi as the interface between CasADi and PyTorch, by running the following command: + pip install l4casadi --no-build-isolation + + +LIMITATIONS: +Since pytorch is wrapped using L4casadi, the casadi functions are generated using External. This means SX variables and +expanding functions are not supported. This will be computationally intensive when solving, making such approach rather slow when +compared to a programmatic approach. Still, it uses much less memory than the symbolic approach, so it has its own advantages. + + + +KNOWN ISSUES: + On Windows (and possibly other platforms), you may randomly get the following error when running this example: + ```python + OMP: Error #15: Initializing libomp.dll, but found libiomp5md.dll already initialized. + ``` + This error comes from the fact that installing the libraries copies the libiomp5md.dll file at two different locations. + When it tries to load them, it notices that "another" library is already loaded. If you are 100% sure that both libraries + are the exact same version, you can safely ignore this error. To do so, you can add the following lines at the beginning + of your script: + ```python + import os + os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" + ``` + That being said, this can be very problematic if the two libraries are not the exact same version. The safer approach is + to delete one of the file. To do so, simply navigate to the folders where the files libiomp5md.dll are located and delete it + (or back-it up by adding ".bak" to the name) keeping only one (keeping the one in site-packages seems to work fine). +""" + +from casadi import MX, vertcat, Function import l4casadi as l4c -import numpy as np import torch -# """ -# INSTALLATION: -# First, make sure pytorch is installed - -# pip install torch>=2.0 --index-url https://download.pytorch.org/whl/cpu/torch_stable.html -# setuptools>=68.1 -# scikit-build>=0.17 -# cmake>=3.27 -# ninja>=1.11 -# Then, install l4casadi as the interface between CasADi and PyTorch - -# pip install l4casadi --no-build-isolation - -# """ - class TorchModel: """ This class wraps a pytorch model and allows the user to call some useful functions on it. """ - def __init__(self, torch_model: torch.nn.Module): - self._dynamic_model = l4c.L4CasADi(torch_model, device="cpu") # device='cuda' for GPU + def __init__(self, torch_model: torch.nn.Module, device="cuda" if torch.cuda.is_available() else "cpu"): + self._dynamic_model = l4c.L4CasADi( + torch_model, device=device, generate_jac_jac=True, generate_adj1=False, generate_jac_adj1=False + ) self._nb_dof = torch_model.size_in // 3 self._symbolic_variables() @@ -66,6 +95,9 @@ def nb_tau(self) -> int: return self.nb_dof def forward_dynamics(self, with_contact: bool = False) -> Function: + if with_contact: + raise NotImplementedError("Contact dynamics are not implemented for torch models") + return Function( "forward_dynamics", [self.q, self.qdot, self.tau, self.external_forces, self.parameters], diff --git a/external/l4casadi b/external/l4casadi new file mode 160000 index 000000000..f01a85d6e --- /dev/null +++ b/external/l4casadi @@ -0,0 +1 @@ +Subproject commit f01a85d6e0151b6178c257aa0e9691219ea05c10 diff --git a/tests/shard3/test_global_getting_started.py b/tests/shard3/test_global_getting_started.py index 899233b66..212d265a8 100644 --- a/tests/shard3/test_global_getting_started.py +++ b/tests/shard3/test_global_getting_started.py @@ -4,6 +4,7 @@ import tracemalloc import gc +import os import pickle import platform import re @@ -2515,3 +2516,16 @@ def test_memory_and_execution_time(): npt.assert_array_less(test_memory[key][0], ref[key][0] * factor) npt.assert_array_less(test_memory[key][1], ref[key][1] * factor) npt.assert_array_less(test_memory[key][2], ref[key][2] * factor) + + +def test_deep_neural_network(): + from bioptim.examples.deep_neural_network import pytorch_ocp as ocp_module + + model = ocp_module.NeuralNetworkModel(layer_node_count=(6, 8, 2), dropout_probability=0.2) + ocp = ocp_module.prepare_ocp(model=model, final_time=1, n_shooting=3) + solver = Solver.IPOPT() + solver.set_maximum_iterations(1) + + # We can launch the solving, but won't be able to check the results as the model is not trained so it is random + os.environ["KMP_DUPLICATE_LIB_OK"] = "True" + ocp.solve(solver=solver) From 9eedac6d2a4f0135239a2005d435d25131018d64 Mon Sep 17 00:00:00 2001 From: Ipuch Date: Thu, 9 Jan 2025 12:13:13 -0500 Subject: [PATCH 10/19] githubaction: miniforge update githubactions: conda githubaction: once for all by mamba --- .github/workflows/run_tests_linux.yml | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/.github/workflows/run_tests_linux.yml b/.github/workflows/run_tests_linux.yml index dfc7f273d..a198349df 100644 --- a/.github/workflows/run_tests_linux.yml +++ b/.github/workflows/run_tests_linux.yml @@ -83,6 +83,19 @@ jobs: - name: Checkout code uses: actions/checkout@v3 + - name: Setup environment + uses: conda-incubator/setup-miniconda@v2 + with: + miniforge-version: latest + use-mamba: true + activate-environment: bioptim + environment-file: environment.yml + + - name: Print mamba info + run: | + conda info +# mamba list + - name: Install extra dependencies run: | sudo apt install -y python3-pip From 4b3b91a824d51c935f04be912ad6e7053a6d82ab Mon Sep 17 00:00:00 2001 From: Pariterre Date: Mon, 13 Jan 2025 16:14:51 -0500 Subject: [PATCH 11/19] Moved installation of torch --- .github/workflows/run_tests_linux.yml | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/.github/workflows/run_tests_linux.yml b/.github/workflows/run_tests_linux.yml index a198349df..a7ed834a7 100644 --- a/.github/workflows/run_tests_linux.yml +++ b/.github/workflows/run_tests_linux.yml @@ -44,13 +44,6 @@ jobs: cd .. if: matrix.shard == 1 - - name: Test installed version of bioptim - run: | - python setup.py install - cd - python -c "import bioptim" - if: matrix.shard == 1 - - name: Install pytorch on Linux run: | pip install torch>=2.0 @@ -59,6 +52,13 @@ jobs: cd ../.. if: matrix.shard == 3 + - name: Test installed version of bioptim + run: | + python setup.py install + cd + python -c "import bioptim" + if: matrix.shard == 1 + - name: Run tests with code coverage run: pytest -v --color=yes --cov-report term-missing --cov=bioptim tests/shard${{ matrix.shard }} --mpl-baseline-path=bioptim/tests/plot_reference_images if: matrix.os == 'ubuntu-latest' From a1e043cebaf502e6550e9e03f5dfffa9fb86420d Mon Sep 17 00:00:00 2001 From: Pariterre Date: Mon, 13 Jan 2025 16:25:45 -0500 Subject: [PATCH 12/19] Added git submodule init --- .github/workflows/run_tests_linux.yml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/run_tests_linux.yml b/.github/workflows/run_tests_linux.yml index a7ed834a7..479184720 100644 --- a/.github/workflows/run_tests_linux.yml +++ b/.github/workflows/run_tests_linux.yml @@ -46,8 +46,10 @@ jobs: - name: Install pytorch on Linux run: | - pip install torch>=2.0 + git submodule update --init --recursive cd external/l4casadi + pip install torch>=2.0 + pip install setuptools>=68.1 scikit-build>=0.17 cmake>=3.27 ninja>=1.11 pip install . --no-build-isolation cd ../.. if: matrix.shard == 3 From da8111d7ff381176a3f3274054650482d04a4290 Mon Sep 17 00:00:00 2001 From: Pariterre Date: Mon, 13 Jan 2025 16:33:06 -0500 Subject: [PATCH 13/19] Added CUDA to github actions --- .github/workflows/run_tests_linux.yml | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/.github/workflows/run_tests_linux.yml b/.github/workflows/run_tests_linux.yml index 479184720..ffa7d36dd 100644 --- a/.github/workflows/run_tests_linux.yml +++ b/.github/workflows/run_tests_linux.yml @@ -44,9 +44,16 @@ jobs: cd .. if: matrix.shard == 1 + - name: Install CUDA toolkit + uses: Jimver/cuda-toolkit@v0.2.19 + id: cuda-toolkit + with: + cuda: '12.5.0' + if: matrix.shard == 3 + - name: Install pytorch on Linux run: | - git submodule update --init --recursive + git submodule update --init --recursive external/l4casadi cd external/l4casadi pip install torch>=2.0 pip install setuptools>=68.1 scikit-build>=0.17 cmake>=3.27 ninja>=1.11 From 33c8c9376ac97229e72e19e39a292e213a5a8bc9 Mon Sep 17 00:00:00 2001 From: Pariterre Date: Mon, 13 Jan 2025 16:44:29 -0500 Subject: [PATCH 14/19] Removed the tests for TorchModel as CUDA can currently be installed on github --- .github/workflows/run_tests_linux.yml | 17 ----------------- 1 file changed, 17 deletions(-) diff --git a/.github/workflows/run_tests_linux.yml b/.github/workflows/run_tests_linux.yml index ffa7d36dd..dcde439fe 100644 --- a/.github/workflows/run_tests_linux.yml +++ b/.github/workflows/run_tests_linux.yml @@ -44,23 +44,6 @@ jobs: cd .. if: matrix.shard == 1 - - name: Install CUDA toolkit - uses: Jimver/cuda-toolkit@v0.2.19 - id: cuda-toolkit - with: - cuda: '12.5.0' - if: matrix.shard == 3 - - - name: Install pytorch on Linux - run: | - git submodule update --init --recursive external/l4casadi - cd external/l4casadi - pip install torch>=2.0 - pip install setuptools>=68.1 scikit-build>=0.17 cmake>=3.27 ninja>=1.11 - pip install . --no-build-isolation - cd ../.. - if: matrix.shard == 3 - - name: Test installed version of bioptim run: | python setup.py install From 79d542b9f95d50c69499b4821f2670960dd4c32b Mon Sep 17 00:00:00 2001 From: Pariterre Date: Mon, 13 Jan 2025 16:46:34 -0500 Subject: [PATCH 15/19] Removed submodule external/l4casadi --- .gitmodules | 5 +---- external/l4casadi | 1 - 2 files changed, 1 insertion(+), 5 deletions(-) delete mode 160000 external/l4casadi diff --git a/.gitmodules b/.gitmodules index 633e50b64..b867fe47b 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,6 +1,3 @@ [submodule "external/acados"] path = external/acados - url = https://github.com/acados/acados.git -[submodule "external/l4casadi"] - path = external/l4casadi - url = https://github.com/pariterre/l4casadi + url = https://github.com/acados/acados.git \ No newline at end of file diff --git a/external/l4casadi b/external/l4casadi deleted file mode 160000 index f01a85d6e..000000000 --- a/external/l4casadi +++ /dev/null @@ -1 +0,0 @@ -Subproject commit f01a85d6e0151b6178c257aa0e9691219ea05c10 From 3f5bfa38ed53a20babaeb057019bbc961b39ba54 Mon Sep 17 00:00:00 2001 From: Pariterre Date: Mon, 27 Jan 2025 14:06:55 -0500 Subject: [PATCH 16/19] Made sure drive is in capital letter on Windows for the tests --- tests/shard1/test_biorbd_model_holonomic.py | 4 +-- tests/shard1/test_prepare_all_examples.py | 1 + tests/shard3/test_global_torque_driven_ocp.py | 29 +++++++++++++++++++ tests/shard3/test_initial_condition.py | 8 ++--- 4 files changed, 36 insertions(+), 6 deletions(-) diff --git a/tests/shard1/test_biorbd_model_holonomic.py b/tests/shard1/test_biorbd_model_holonomic.py index 490fcbe65..332b2ae60 100644 --- a/tests/shard1/test_biorbd_model_holonomic.py +++ b/tests/shard1/test_biorbd_model_holonomic.py @@ -1,10 +1,10 @@ import platform +from bioptim import HolonomicBiorbdModel, HolonomicConstraintsFcn, HolonomicConstraintsList, Solver, SolutionMerge +from casadi import DM, MX import numpy as np import numpy.testing as npt import pytest -from casadi import DM, MX -from bioptim import HolonomicBiorbdModel, HolonomicConstraintsFcn, HolonomicConstraintsList, Solver, SolutionMerge from ..utils import TestUtils diff --git a/tests/shard1/test_prepare_all_examples.py b/tests/shard1/test_prepare_all_examples.py index 4704ec875..deb987974 100644 --- a/tests/shard1/test_prepare_all_examples.py +++ b/tests/shard1/test_prepare_all_examples.py @@ -1,3 +1,4 @@ +from bioptim import InterpolationType, PhaseDynamics, OdeSolver import numpy as np import pytest diff --git a/tests/shard3/test_global_torque_driven_ocp.py b/tests/shard3/test_global_torque_driven_ocp.py index 1f8dd6cda..f633ac572 100644 --- a/tests/shard3/test_global_torque_driven_ocp.py +++ b/tests/shard3/test_global_torque_driven_ocp.py @@ -283,6 +283,35 @@ def test_track_marker_2D_pendulum(ode_solver, phase_dynamics): if platform.system() == "Windows": return + bioptim_folder = TestUtils.module_folder(ocp_module) + + ode_solver_orig = ode_solver + ode_solver = ode_solver() + + # Define the problem + model_path = bioptim_folder + "/models/pendulum.bioMod" + bio_model = TorqueBiorbdModel(model_path) + + final_time = 2 + n_shooting = 30 + + # Generate data to fit + np.random.seed(42) + markers_ref = np.random.rand(3, 2, n_shooting + 1) + tau_ref = np.random.rand(2, n_shooting) + + if isinstance(ode_solver, OdeSolver.IRK): + tau_ref = tau_ref * 5 + + ocp = ocp_module.prepare_ocp( + bio_model, + final_time, + n_shooting, + markers_ref, + tau_ref, + ode_solver=ode_solver, + expand_dynamics=ode_solver_orig != OdeSolver.IRK, + ) sol = ocp.solve() # Check constraints diff --git a/tests/shard3/test_initial_condition.py b/tests/shard3/test_initial_condition.py index 5a12d0cb3..989634a47 100644 --- a/tests/shard3/test_initial_condition.py +++ b/tests/shard3/test_initial_condition.py @@ -1,9 +1,5 @@ import re -import numpy as np -import numpy.testing as npt -import pytest - from bioptim import ( InterpolationType, Solution, @@ -19,6 +15,10 @@ SolutionMerge, ) from bioptim.limits.path_conditions import InitialGuess +import numpy as np +import numpy.testing as npt +import pytest + from ..utils import TestUtils From 1a2f997357c5fd92db62eb3f56613ba7012ef475 Mon Sep 17 00:00:00 2001 From: Pariterre Date: Wed, 26 Feb 2025 16:47:35 -0500 Subject: [PATCH 17/19] Updated setup-miniconda to v3 --- .github/workflows/run_tests_linux.yml | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/.github/workflows/run_tests_linux.yml b/.github/workflows/run_tests_linux.yml index dcde439fe..b1dce64ff 100644 --- a/.github/workflows/run_tests_linux.yml +++ b/.github/workflows/run_tests_linux.yml @@ -76,17 +76,16 @@ jobs: uses: actions/checkout@v3 - name: Setup environment - uses: conda-incubator/setup-miniconda@v2 + uses: conda-incubator/setup-miniconda@v3 with: miniforge-version: latest - use-mamba: true activate-environment: bioptim environment-file: environment.yml - - name: Print mamba info + - name: Print conda info run: | conda info -# mamba list + conda list - name: Install extra dependencies run: | From b52271b938b49c0752be46568d927bdd16f7ed2d Mon Sep 17 00:00:00 2001 From: Pariterre Date: Tue, 9 Sep 2025 10:23:38 -0400 Subject: [PATCH 18/19] Fixed some mischanges from the rebase --- tests/shard1/test_prepare_all_examples.py | 1 - tests/shard3/test_global_torque_driven_ocp.py | 29 ------------------- 2 files changed, 30 deletions(-) diff --git a/tests/shard1/test_prepare_all_examples.py b/tests/shard1/test_prepare_all_examples.py index deb987974..4704ec875 100644 --- a/tests/shard1/test_prepare_all_examples.py +++ b/tests/shard1/test_prepare_all_examples.py @@ -1,4 +1,3 @@ -from bioptim import InterpolationType, PhaseDynamics, OdeSolver import numpy as np import pytest diff --git a/tests/shard3/test_global_torque_driven_ocp.py b/tests/shard3/test_global_torque_driven_ocp.py index f633ac572..1f8dd6cda 100644 --- a/tests/shard3/test_global_torque_driven_ocp.py +++ b/tests/shard3/test_global_torque_driven_ocp.py @@ -283,35 +283,6 @@ def test_track_marker_2D_pendulum(ode_solver, phase_dynamics): if platform.system() == "Windows": return - bioptim_folder = TestUtils.module_folder(ocp_module) - - ode_solver_orig = ode_solver - ode_solver = ode_solver() - - # Define the problem - model_path = bioptim_folder + "/models/pendulum.bioMod" - bio_model = TorqueBiorbdModel(model_path) - - final_time = 2 - n_shooting = 30 - - # Generate data to fit - np.random.seed(42) - markers_ref = np.random.rand(3, 2, n_shooting + 1) - tau_ref = np.random.rand(2, n_shooting) - - if isinstance(ode_solver, OdeSolver.IRK): - tau_ref = tau_ref * 5 - - ocp = ocp_module.prepare_ocp( - bio_model, - final_time, - n_shooting, - markers_ref, - tau_ref, - ode_solver=ode_solver, - expand_dynamics=ode_solver_orig != OdeSolver.IRK, - ) sol = ocp.solve() # Check constraints From 808dfa6053db031149d624f9854e02e63ba80dff Mon Sep 17 00:00:00 2001 From: Pariterre Date: Thu, 25 Sep 2025 16:40:38 -0400 Subject: [PATCH 19/19] Moved the example into toy_example --- .../{ => toy_examples}/deep_neural_network/pytorch_ocp.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename bioptim/examples/{ => toy_examples}/deep_neural_network/pytorch_ocp.py (100%) diff --git a/bioptim/examples/deep_neural_network/pytorch_ocp.py b/bioptim/examples/toy_examples/deep_neural_network/pytorch_ocp.py similarity index 100% rename from bioptim/examples/deep_neural_network/pytorch_ocp.py rename to bioptim/examples/toy_examples/deep_neural_network/pytorch_ocp.py