diff --git a/bioptim/examples/getting_started/_l4c_generated/adj1_l4casadi_f.pt b/bioptim/examples/getting_started/_l4c_generated/adj1_l4casadi_f.pt new file mode 100644 index 000000000..63d2c1b8d Binary files /dev/null and b/bioptim/examples/getting_started/_l4c_generated/adj1_l4casadi_f.pt differ diff --git a/bioptim/examples/getting_started/_l4c_generated/jac_adj1_l4casadi_f.pt b/bioptim/examples/getting_started/_l4c_generated/jac_adj1_l4casadi_f.pt new file mode 100644 index 000000000..96fc6b382 Binary files /dev/null and b/bioptim/examples/getting_started/_l4c_generated/jac_adj1_l4casadi_f.pt differ diff --git a/bioptim/examples/getting_started/_l4c_generated/jac_l4casadi_f.pt b/bioptim/examples/getting_started/_l4c_generated/jac_l4casadi_f.pt new file mode 100644 index 000000000..5c47ddff2 Binary files /dev/null and b/bioptim/examples/getting_started/_l4c_generated/jac_l4casadi_f.pt differ diff --git a/bioptim/examples/getting_started/_l4c_generated/l4casadi_f.cpp b/bioptim/examples/getting_started/_l4c_generated/l4casadi_f.cpp new file mode 100644 index 000000000..730d5df7f --- /dev/null +++ b/bioptim/examples/getting_started/_l4c_generated/l4casadi_f.cpp @@ -0,0 +1,111 @@ +#include + +L4CasADi l4casadi("/home/pariterre/Programmation/Technopole/bioptim/bioptim/examples/getting_started/_l4c_generated", "l4casadi_f", 1, 6, 1, 2, "cpu", true, true, true, false, true, false); + +#ifdef __cplusplus +extern "C" { +#endif + +#include + +#ifndef casadi_real +#define casadi_real double +#endif + +#ifndef casadi_int +#define casadi_int long long int +#endif + +/* Symbol visibility in DLLs */ +#ifndef CASADI_SYMBOL_EXPORT + #if defined(_WIN32) || defined(__WIN32__) || defined(__CYGWIN__) + #if defined(STATIC_LINKED) + #define CASADI_SYMBOL_EXPORT + #else + #define CASADI_SYMBOL_EXPORT __declspec(dllexport) + #endif + #elif defined(__GNUC__) && defined(GCC_HASCLASSVISIBILITY) + #define CASADI_SYMBOL_EXPORT __attribute__ ((visibility ("default"))) + #else + #define CASADI_SYMBOL_EXPORT + #endif +#endif + +// Function l4casadi_f + +static const casadi_int l4casadi_f_s_in0[3] = { 1, 6, 1}; +static const casadi_int l4casadi_f_s_out0[3] = { 1, 2, 1}; + +// Only single input, single output is supported at the moment +CASADI_SYMBOL_EXPORT casadi_int l4casadi_f_n_in(void) { return 1;} +CASADI_SYMBOL_EXPORT casadi_int l4casadi_f_n_out(void) { return 1;} + +CASADI_SYMBOL_EXPORT const casadi_int* l4casadi_f_sparsity_in(casadi_int i) { + switch (i) { + case 0: return l4casadi_f_s_in0; + default: return 0; + } +} + +CASADI_SYMBOL_EXPORT const casadi_int* l4casadi_f_sparsity_out(casadi_int i) { + switch (i) { + case 0: return l4casadi_f_s_out0; + default: return 0; + } +} + +CASADI_SYMBOL_EXPORT int l4casadi_f(const casadi_real** arg, casadi_real** res, casadi_int* iw, casadi_real* w, int mem){ + l4casadi.forward(arg[0], res[0]); + return 0; +} + +// Jacobian l4casadi_f + +CASADI_SYMBOL_EXPORT casadi_int jac_l4casadi_f_n_in(void) { return 2;} +CASADI_SYMBOL_EXPORT casadi_int jac_l4casadi_f_n_out(void) { return 1;} + +CASADI_SYMBOL_EXPORT int jac_l4casadi_f(const casadi_real** arg, casadi_real** res, casadi_int* iw, casadi_real* w, int mem){ + l4casadi.jac(arg[0], res[0]); + return 0; +} + + + +// adj1 l4casadi_f + +CASADI_SYMBOL_EXPORT casadi_int adj1_l4casadi_f_n_in(void) { return 3;} +CASADI_SYMBOL_EXPORT casadi_int adj1_l4casadi_f_n_out(void) { return 1;} + +CASADI_SYMBOL_EXPORT int adj1_l4casadi_f(const casadi_real** arg, casadi_real** res, casadi_int* iw, casadi_real* w, int mem){ + // adj1 [i0, out_o0, adj_o0] -> [out_adj_i0] + l4casadi.adj1(arg[0], arg[2], res[0]); + return 0; +} + + +// jac_adj1 l4casadi_f + +CASADI_SYMBOL_EXPORT casadi_int jac_adj1_l4casadi_f_n_in(void) { return 4;} +CASADI_SYMBOL_EXPORT casadi_int jac_adj1_l4casadi_f_n_out(void) { return 3;} + +CASADI_SYMBOL_EXPORT int jac_adj1_l4casadi_f(const casadi_real** arg, casadi_real** res, casadi_int* iw, casadi_real* w, int mem){ + // jac_adj1 [i0, out_o0, adj_o0, out_adj_i0] -> [jac_adj_i0_i0, jac_adj_i0_out_o0, jac_adj_i0_adj_o0] + if (res[1] != NULL) { + l4casadi.invalid_argument("jac_adj_i0_out_o0 is not provided by L4CasADi. If you need this feature, please contact the L4CasADi developer."); + } + if (res[2] != NULL) { + l4casadi.invalid_argument("jac_adj_i0_adj_o0 is not provided by L4CasADi. If you need this feature, please contact the L4CasADi developer."); + } + if (res[0] == NULL) { + l4casadi.invalid_argument("L4CasADi can only provide jac_adj_i0_i0 for jac_adj1_l4casadi_f function. If you need this feature, please contact the L4CasADi developer."); + } + l4casadi.jac_adj1(arg[0], arg[2], res[0]); + return 0; +} + + + + +#ifdef __cplusplus +} /* extern "C" */ +#endif \ No newline at end of file diff --git a/bioptim/examples/getting_started/_l4c_generated/l4casadi_f.pt b/bioptim/examples/getting_started/_l4c_generated/l4casadi_f.pt new file mode 100644 index 000000000..8c503c1ff Binary files /dev/null and b/bioptim/examples/getting_started/_l4c_generated/l4casadi_f.pt differ 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()