diff --git a/.gitignore b/.gitignore index 527bde6..aee7883 100644 --- a/.gitignore +++ b/.gitignore @@ -135,3 +135,5 @@ dmypy.json # DS_STORE .DS_Store +*.mp4 +*.gif \ No newline at end of file diff --git a/examples/benchmark_planar_pcs.py b/examples/benchmark_planar_pcs.py new file mode 100644 index 0000000..759f673 --- /dev/null +++ b/examples/benchmark_planar_pcs.py @@ -0,0 +1,1594 @@ +import cv2 # importing cv2 +from functools import partial +import jax + +jax.config.update("jax_enable_x64", True) # double precision +from diffrax import diffeqsolve, Euler, ODETerm, SaveAt, Tsit5 +from jax import Array, vmap, jit +from jax import numpy as jnp +import matplotlib.pyplot as plt +import numpy as onp +from pathlib import Path +from typing import Callable, Dict, Optional, Literal, Union + +import jsrm +from jsrm import ode_factory +from jsrm.systems import planar_pcs_num, planar_pcs_sym + +import time +import pickle + +import timeit +import statistics + +def simulate_planar_pcs_value_eval( + num_segments: int, + type_of_derivation: Optional[Literal["symbolic", "numeric"]] = "symbolic", + type_of_integration: Optional[Literal["gauss-legendre", "gauss-kronrad", "trapezoid"]] = "gauss-legendre", + param_integration: Optional[int] = None, + type_of_jacobian: Optional[Literal["explicit", "autodiff"]] = "explicit", + robot_params: Optional[Dict[str, Array]] = None, + strain_selector: Optional[Array] = None, + q0: Optional[Array] = None, + q_d0: Optional[Array] = None, + t: float = 1.0, + dt: Optional[float] = None, + bool_print: bool = True, + bool_plot: bool = True, + bool_save_plot: bool = False, + bool_save_video: bool = True, + bool_save_res: bool = False, + results_path: Optional[Union[str, Path]] = None, + results_path_extension: Optional[str] = None +) -> Dict: + """ + Simulate a planar PCS model. Save the video and figures. + + Args: + num_segments (int): number of segments of the robot. + type_of_derivation (str, optional): type of derivation ("symbolic" or "numeric"). + Defaults to "symbolic". + type_of_integration (str, optional): scheme of integration ("gauss-legendre", "gauss-kronrad" or "trapezoid"). + Defaults to "gauss-legendre". + param_integration (int, optional): parameter for the integration scheme (number of points for gauss or number of points for trapezoid). + Defaults to 1000 for trapezoid and 5 for gauss. + type_of_jacobian (str, optional): type of jacobian ("explicit" or "autodiff"). + Defaults to "explicit". + strain_selector (Array, optional): selector for the strains as a boolean array. + Defaults to all. + robot_params (Dict[str, Array], optional): parameters of the robot. + strain_selector (Array, optional): selector for the strains. + t (float, optional): time of simulation [s]. + Defaults to 1.0s. + dt (float, optional): time step [s]. + Defaults to 1e-4s. + q0 (Array, optional): initial configuration. + Defaults to a configuration with 5.0pi rad of bending, 0.1 of shear, and 0.2 of axial strain for each segment. + q_d0 (Array, optional): initial velocity. + Defaults to 0 array. + bool_print (bool, optional): if True, print the simulation results. + Defaults to True. + bool_plot (bool, optional): if True, show the figures. + Defaults to True. + bool_save_plot (bool, optional): if True, save the figures. + Defaults to False. + bool_save_video (bool, optional): if True, save the video. + Defaults to False. + bool_save_res (bool, optional): if True, save the results of the simulation. + Defaults to False. + results_path (str, optional): path to save the dictionary with the simulation results. Must have the suffix .pkl. + results_path_extension (str, optional): extension to add to the results path. + Defaults to None, which will use the default path. + + Returns: + simulation_dict (Dict): dictionary with the simulation results. + """ + + # =================================================== + # Initialization of the simulation parameters + # =================================================== + + if not isinstance(num_segments, int): + raise TypeError(f"num_segments must be an integer, but got {type(num_segments).__name__}") + if num_segments < 1: + raise ValueError(f"num_segments must be greater than 0, but got {num_segments}") + + if not isinstance(type_of_derivation, str): + raise TypeError(f"type_of_derivation must be a string, but got {type(type_of_derivation).__name__}") + if type_of_derivation == "numeric": + if not isinstance(type_of_integration, str): + raise TypeError(f"type_of_integration must be a string, but got {type(type_of_integration).__name__}") + if param_integration is None: + if type_of_integration == "gauss-legendre": + param_integration = 5 + if type_of_integration == "gauss-kronrad": + param_integration = 15 + elif type_of_integration == "trapezoid": + param_integration = 1000 + if not isinstance(param_integration, int): + raise TypeError(f"param_integration must be an integer, but got {type(param_integration).__name__}") + if param_integration < 1: + raise ValueError(f"param_integration must be greater than 0, but got {param_integration}") + + if type_of_jacobian not in ["explicit", "autodiff"]: + raise ValueError( + f"type_of_jacobian must be 'explicit' or 'autodiff', but got {type_of_jacobian}" + ) + + elif type_of_derivation == "symbolic": + # filepath to symbolic expressions + sym_exp_filepath = ( + Path(jsrm.__file__).parent + / "symbolic_expressions" + / f"planar_pcs_ns-{num_segments}.dill" + ) + if not sym_exp_filepath.exists(): + raise FileNotFoundError( + f"File {sym_exp_filepath} does not exist. Please run the script to generate the symbolic expressions." + ) + else: + raise ValueError( + f"type_of_derivation must be 'symbolic' or 'numeric', but got {type_of_derivation}" + ) + + if robot_params is None: + # set parameters + rho = 1070 * jnp.ones((num_segments,)) # Volumetric density of Dragon Skin 20 [kg/m^3] + robot_params = { + "th0": jnp.array(0.0), # Initial orientation angle [rad] + "l": 1e-1 * jnp.ones((num_segments,)), # Length of each segment [m] + "r": 2e-2 * jnp.ones((num_segments,)), # Radius of each segment [m] + "rho": rho, + "g": jnp.array([0.0, 9.81]), # Gravity vector [m/s^2] + "E": 2e3 * jnp.ones((num_segments,)), # Elastic modulus [Pa] + "G": 1e3 * jnp.ones((num_segments,)), # Shear modulus [Pa] + } + robot_params["D"] = 1e-3 * jnp.diag( # Damping matrix [Ns/m] + (jnp.repeat( + jnp.array([[1e0, 1e3, 1e3]]), num_segments, axis=0 + ) * robot_params["l"][:, None]).flatten() + ) + #TODO: check if the parameters are correctly defined + + # Max number of degrees of freedom = size of the strain vector + n_xi = 3 * num_segments + + if strain_selector is None: + # activate all strains (i.e. bending, shear, and axial) + strain_selector = jnp.ones((n_xi,), dtype=bool) + if not isinstance(strain_selector, jnp.ndarray): + if isinstance(strain_selector, list): + strain_selector = jnp.array(strain_selector) + else: + raise TypeError(f"strain_selector must be a jnp.ndarray, but got {type(strain_selector).__name__}") + strain_selector = strain_selector.flatten() + if strain_selector.shape[0] != n_xi: + raise ValueError( + f"strain_selector must have the same shape as the strain vector, but got {strain_selector.shape[0]} instead of {n_xi}" + ) + if not jnp.issubdtype(strain_selector.dtype, jnp.bool_): + raise TypeError( + f"strain_selector must be a boolean array, but got {strain_selector.dtype}" + ) + + # number of generalized coordinates + n_dof = jnp.sum(strain_selector) + + if q0 is None: + q0_init = jnp.repeat(jnp.array([5.0 * jnp.pi, 0.1, 0.2])[None, :], num_segments, axis=0).flatten() + q0 = jnp.array([q0_init[i] for i in range(3*num_segments) if strain_selector[i]]) + if not isinstance(q0, jnp.ndarray): + if isinstance(q0, list): + q0 = jnp.array(q0) + else: + raise TypeError(f"q0 must be a jnp.ndarray, but got {type(q0).__name__}") + if q0.shape[0] != n_dof: + raise ValueError( + f"q0 must have the same shape as jnp.sum(strain_selector), but got {q0.shape[0]} instead of {jnp.sum(strain_selector)}" + ) + + if q_d0 is None: + # define initial velocity + q_d0 = jnp.zeros_like(q0) + if not isinstance(q_d0, jnp.ndarray): + if isinstance(q_d0, list): + q_d0 = jnp.array(q_d0) + else: + raise TypeError(f"q_d0 must be a jnp.ndarray, but got {type(q_d0).__name__}") + if q_d0.shape[0] != n_dof: + raise ValueError( + f"q_d0 must have the same shape as q0, but got {q_d0.shape[0]} instead of {n_dof}" + ) + + if not isinstance(t, float): + if isinstance(t, int): + t = float(t) + else: + raise TypeError(f"t must be a float, but got {type(t).__name__}") + if t <= 0: + raise ValueError(f"t must be greater than 0, but got {t}") + + if dt is None: + dt = 1e-4 + if not isinstance(dt, float): + if isinstance(dt, int): + dt = float(dt) + else: + raise TypeError(f"dt must be a float, but got {type(dt).__name__}") + if dt <= 0: + raise ValueError(f"dt must be greater than 0, but got {dt}") + if dt > t: + raise ValueError(f"dt must be less than t, but got {dt} > {t}") + + if not isinstance(bool_print, bool): + raise TypeError(f"bool_print must be a boolean, but got {type(bool_print).__name__}") + if not isinstance(bool_plot, bool): + raise TypeError(f"bool_plot must be a boolean, but got {type(bool_plot).__name__}") + if not isinstance(bool_save_plot, bool): + raise TypeError(f"bool_save_plot must be a boolean, but got {type(bool_save_plot).__name__}") + if not isinstance(bool_save_video, bool): + raise TypeError(f"bool_save_video must be a boolean, but got {type(bool_save_video).__name__}") + if not isinstance(bool_save_res, bool): + raise TypeError(f"bool_save_res must be a boolean, but got {type(bool_save_res).__name__}") + + if bool_save_res: + if results_path is None: + results_path_parent = ( + Path(__file__).parent + / "results" + / "values" + / "planar_pcs_results" + / f"ns-{num_segments}") + file_name = f"{('symb' if type_of_derivation == 'symbolic' else 'num')}" + if type_of_derivation == "numeric": + file_name += f"-{type_of_integration}-{param_integration}-{type_of_jacobian}" + + if results_path_extension is not None: + if not isinstance(results_path_extension, str): + raise TypeError( + f"results_path_extension must be a string, but got {type(results_path_extension).__name__}" + ) + file_name += f"-{results_path_extension}" + + results_path = (results_path_parent / file_name).with_suffix(".pkl") + + if isinstance(results_path, str) or isinstance(results_path, Path): + results_path_obj = Path(results_path) + if results_path_obj.suffix != ".pkl": + raise ValueError( + f"results_path must have the suffix .pkl, but got {results_path_obj.suffix}" + ) + results_path = results_path_obj + else: + raise TypeError( + f"results_path must be a string, but got {type(results_path).__name__}" + ) + + # create the directory if it does not exist + if not results_path.parent.exists(): + results_path.parent.mkdir(parents=True, exist_ok=True) + + # ============================================================================================= + # Figures and video generation + # ============================================================================================= + extension = f"planar_pcs_ns-{num_segments}-{('sym' if type_of_derivation == 'symbolic' else 'num')}" + if type_of_derivation == "numeric": + extension += f"-{type_of_integration}-{param_integration}" + + # ====================== + # For video generation + # ====================== + # set simulation parameters + ts = jnp.arange(0.0, t, dt) # time steps + skip_step = 10 # how many time steps to skip in between video frames + video_ts = ts[::skip_step] # time steps for video + + # video settings + video_width, video_height = 700, 700 # img height and width" + video_path_parent = Path(__file__).parent / "videos" / "planar_pcs" / f"ns-{num_segments}" + if bool_save_video: + video_path_parent.mkdir(parents=True, exist_ok=True) + + video_path = video_path_parent / f"{extension}.mp4" + + def draw_robot( + batched_forward_kinematics_fn: Callable, + params: Dict[str, Array], + q: Array, + width: int, + height: int, + num_points: int = 50, + ) -> onp.ndarray: + """ + Draw the robot in an array of shape (width, height, 3) in OpenCV format. + The robot is drawn as a curve in the x-y plane. The base is drawn as a rectangle. + + Args: + batched_forward_kinematics_fn (Callable): function to compute the forward kinematics. + params (Dict[str, Array]): parameters of the robot. + q (Array): configuration of the robot. + width (int): width of the image. + height (int): height of the image. + num_points (int, optional): number of points to draw the robot. + Defaults to 50. + + Returns: + onp.ndarray: image of the robot in OpenCV format. + """ + # plotting in OpenCV + h, w = height, width # img height and width + ppm = h / (2.0 * jnp.sum(params["l"])) # pixel per meter + base_color = (0, 0, 0) # black robot_color in BGR + robot_color = (255, 0, 0) # black robot_color in BGR + + # we use for plotting N points along the length of the robot + s_ps = jnp.linspace(0, jnp.sum(params["l"]), num_points) + + # poses along the robot of shape (3, N) + chi_ps = batched_forward_kinematics_fn(params, q, s_ps) + + img = 255 * onp.ones((w, h, 3), dtype=jnp.uint8) # initialize background to white + curve_origin = onp.array( + [w // 2, 0.1 * h], dtype=onp.int32 + ) # in x-y pixel coordinates + # draw base + cv2.rectangle(img, (0, h - curve_origin[1]), (w, h), color=base_color, thickness=-1) + # transform robot poses to pixel coordinates + # should be of shape (N, 2) + curve = onp.array((curve_origin + chi_ps[:2, :].T * ppm), dtype=onp.int32) + # invert the v pixel coordinate + curve[:, 1] = h - curve[:, 1] + cv2.polylines(img, [curve], isClosed=False, color=robot_color, thickness=10) + + return img + + # ====================== + # For figure saving + # ====================== + figures_path_parent = Path(__file__).parent / "figures" / "planar_pcs" / f"ns-{num_segments}" + figures_path_parent.mkdir(parents=True, exist_ok=True) + + # ===================================================================================================== + # Simulation + # ===================================================================================================== + # save the simulation parameters and results + simulation_dict = { + "params": { + "num_segments": num_segments, + "type_of_derivation": type_of_derivation, + "type_of_integration": type_of_integration, + "param_integration": param_integration, + "type_of_jacobian": type_of_jacobian, + "robot_params": robot_params, + "strain_selector": strain_selector, + "q0": q0, + "q_d0": q_d0, + "t": t, + "dt": dt, + "video_ts": video_ts + }, + "results": { + }, + "execution_time": { + } + } + if bool_save_video: + simulation_dict["execution_time"]["draw_robot"] = None + + + print("Number of segments:", num_segments) + print("Type of derivation:", type_of_derivation) + if type_of_derivation == "numeric": + print("Type of integration:", type_of_integration) + print("Parameter for integration:", param_integration) + print("Type of jacobian:", type_of_jacobian) + print() + + # ==================================================== + print("Importing the planar PCS model...") + + if type_of_derivation == "symbolic": + strain_basis, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns = ( + planar_pcs_sym.factory(sym_exp_filepath, strain_selector) + ) + elif type_of_derivation == "numeric": + strain_basis, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns = ( + planar_pcs_num.factory( + num_segments, + strain_selector, + integration_type=type_of_integration, + param_integration=param_integration, + jacobian_type=type_of_jacobian + ) + ) + + # ==================================================== + # Dynamical matrices function + + print("JIT-compiling the dynamical matrices function...") + dynamical_matrices_fn = partial(dynamical_matrices_fn) + dynamical_matrices_fn = jit(dynamical_matrices_fn) + + print("Computing the dynamical matrices for the initial configuration and velocity...") + B, C, G, K, D, A = dynamical_matrices_fn(robot_params, q0, q_d0) + + simulation_dict["results"]["B_0"] = B + simulation_dict["results"]["C_0"] = C + simulation_dict["results"]["G_0"] = G + simulation_dict["results"]["K_0"] = K + simulation_dict["results"]["D_0"] = D + simulation_dict["results"]["A_0"] = A + + # ==================================================== + # ODE function + + # Parameter for the simulation + x0 = jnp.concatenate([q0, q_d0]) # initial condition + tau = jnp.zeros_like(q0) # torques + + print("Creating the ODE function...") + ode_fn = ode_factory(dynamical_matrices_fn, robot_params, tau) + term = ODETerm(ode_fn) + + print("JIT-compiling the ODE function...") + diffeqsolve_fn = partial(diffeqsolve, + term, + solver=Tsit5(), + t0=ts[0], + t1=ts[-1], + dt0=dt, + y0=x0, + max_steps=None, + saveat=SaveAt(ts=video_ts)) + diffeqsolve_fn = jit(diffeqsolve_fn) # type: ignore + + print("Solving the ODE ...") + sol = diffeqsolve_fn() + q_ts = sol.ys[:, :n_dof] # the evolution of the generalized coordinates + q_d_ts = sol.ys[:, n_dof:] # the evolution of the generalized velocities + + if bool_print: + print("q_ts =\n", q_ts) + print("q_d_ts =\n", q_d_ts) + + simulation_dict["results"]["q_ts"] = q_ts + simulation_dict["results"]["q_d_ts"] = q_d_ts + + # ==================================================== + # Forward kinematics function + s_max = jnp.array([jnp.sum(robot_params["l"])]) # maximum length of the robot + + + forward_kinematics_fn_end_effector = partial(forward_kinematics_fn, robot_params, s=s_max) + + print("JIT-compiling the forward kinematics function...") + forward_kinematics_fn_end_effector = jit(forward_kinematics_fn_end_effector) # type: ignore + forward_kinematics_fn_end_effector = vmap(forward_kinematics_fn_end_effector) + + print("Computing the end-effector position along the trajectory...") + chi_ee_ts = forward_kinematics_fn_end_effector(q_ts) + + if bool_print: + print("chi_ee_ts =\n", chi_ee_ts) + + simulation_dict["results"]["chi_ee_ts"] = chi_ee_ts + + #==================================================== + # Plotting + #=================== + if bool_plot or bool_save_plot: + print("Plotting the results... \n") + + + # Plot the configuration vs time + plt.figure() + plt.title("Configuration vs Time") + for segment_idx in range(num_segments): + plt.plot( + video_ts, q_ts[:, 3 * segment_idx + 0], + label=r"$\kappa_\mathrm{be," + str(segment_idx + 1) + "}$ [rad/m]" + ) + plt.plot( + video_ts, q_ts[:, 3 * segment_idx + 1], + label=r"$\sigma_\mathrm{sh," + str(segment_idx + 1) + "}$ [-]" + ) + plt.plot( + video_ts, q_ts[:, 3 * segment_idx + 2], + label=r"$\sigma_\mathrm{ax," + str(segment_idx + 1) + "}$ [-]" + ) + plt.xlabel("Time [s]") + plt.ylabel("Configuration") + plt.legend() + plt.grid(True) + plt.tight_layout() + if bool_save_plot: + plt.savefig( + figures_path_parent / f"config_vs_time_{extension}.png", bbox_inches="tight", dpi=300 + ) + print("Figures saved at", figures_path_parent / f"config_vs_time_{extension}.png") + if bool_plot: + plt.show() + plt.close() + + # Plot end-effector position vs time + plt.figure() + plt.title("End-effector position vs Time") + plt.plot(video_ts, chi_ee_ts[:, 0], label="x") + plt.plot(video_ts, chi_ee_ts[:, 1], label="y") + plt.xlabel("Time [s]") + plt.ylabel("End-effector Position [m]") + plt.legend() + plt.grid(True) + plt.box(True) + plt.tight_layout() + if bool_save_plot: + plt.savefig( + figures_path_parent / f"end_effector_position_vs_time_{extension}.png", bbox_inches="tight", dpi=300 + ) + print("Figures saved at", figures_path_parent / f"end_effector_position_vs_time_{extension}.png ") + if bool_plot: + plt.show() + plt.close() + + # plot the end-effector position in the x-y plane as a scatter plot with the time as the color + plt.figure() + plt.title("End-effector position in the x-y plane") + plt.scatter(chi_ee_ts[:, 0], chi_ee_ts[:, 1], c=video_ts, cmap="viridis") + plt.axis("equal") + plt.grid(True) + plt.xlabel("End-effector x [m]") + plt.ylabel("End-effector y [m]") + plt.colorbar(label="Time [s]") + plt.tight_layout() + if bool_save_plot: + plt.savefig( + figures_path_parent / f"end_effector_position_xy_{extension}.png", bbox_inches="tight", dpi=300 + ) + print("Figures saved at", figures_path_parent / f"end_effector_position_xy_{extension}.png \n") + if bool_plot: + plt.show() + plt.close() + + # ==================================================== + # Potential energy + + potential_energy_fn = auxiliary_fns["potential_energy_fn"] + potential_energy_fn = partial(potential_energy_fn, robot_params) + + print("JIT-compiling the potential energy function...") + potential_energy_fn = jit(potential_energy_fn) + potential_energy_fn = vmap(potential_energy_fn) + + print("Computing the potential energy along the trajectory...") + U_ts = potential_energy_fn(q_ts) + + simulation_dict["results"]["U_ts"] = U_ts + + # ==================================================== + # Kinetic energy + + kinetic_energy_fn = auxiliary_fns["kinetic_energy_fn"] + kinetic_energy_fn = partial(kinetic_energy_fn, robot_params) + + print("JIT-compiling the kinetic energy function...") + kinetic_energy_fn = jit(kinetic_energy_fn) + kinetic_energy_fn = vmap(kinetic_energy_fn) + + print("Computing the kinetic energy along the trajectory...") + T_ts = kinetic_energy_fn(q_ts, q_d_ts) + + simulation_dict["results"]["T_ts"] = T_ts + + if bool_print: + print("U_ts =\n", U_ts) + print("T_ts =\n", T_ts) + + if bool_plot or bool_save_plot: + print("Plotting the energy... \n") + # Plot the energy vs time + plt.figure() + plt.title("Energy vs Time") + plt.plot(video_ts, U_ts, label="Potential energy") + plt.plot(video_ts, T_ts, label="Kinetic energy") + plt.xlabel("Time [s]") + plt.ylabel("Energy [J]") + plt.legend() + plt.grid(True) + plt.box(True) + plt.tight_layout() + if bool_save_plot: + plt.savefig( + figures_path_parent / f"energy_vs_time_{extension}.png", bbox_inches="tight", dpi=300 + ) + print("Figures saved at", figures_path_parent / f"energy_vs_time_{extension}.png \n") + if bool_plot: + plt.show() + plt.close() + + # ==================================================== + # Video generation + # ================= + if bool_save_video: + print("Drawing the robot...") + + # Vectorize the forward kinematics function according to the s coordinates + print("Vectorizing the forward kinematics function according to the number of segments...") + batched_forward_kinematics = vmap( + forward_kinematics_fn, in_axes=(None, None, 0), out_axes=-1 + ) + + # Initialize the video + fourcc = cv2.VideoWriter_fourcc(*"mp4v") # type: ignore + video = cv2.VideoWriter( + str(video_path), + fourcc, + 1 / (skip_step * dt), # fps + (video_width, video_height), + ) + + # Create the video frames + for time_idx, t in enumerate(video_ts): + x = sol.ys[time_idx] + img = draw_robot( + batched_forward_kinematics, + robot_params, + x[: (x.shape[0] // 2)], + video_width, + video_height, + ) + video.write(img) + + # Release the video + video.release() + + print(f"Video saved at {video_path}. \n") + + # =========================================================================== + # Save the simulation results + # =========================== + if bool_save_res: + print("Saving the simulation results...") + assert results_path is not None, "results_path should not be None when saving results" + with open(results_path, "wb") as f: + pickle.dump(simulation_dict, f) + print(f"Simulation results saved at {results_path} \n") + else: + print("Simulation results not saved. \n") + + print("Simulation finished. \n") + print("**************************************************************** \n \n") + + return simulation_dict + +EVAL_WITHOUT_JIT = False +def simulate_planar_pcs_time_eval( + num_segments: int, + type_of_derivation: str = "symbolic", + type_of_integration: str = "gauss-legendre", + param_integration: Optional[int] = None, + type_of_jacobian: str = "explicit", + robot_params: Optional[Dict[str, Array]] = None, + strain_selector: Optional[Array] = None, + q0: Optional[Array] = None, + q_d0: Optional[Array] = None, + t: float = 1.0, + dt: Optional[float] = None, + bool_save_res: bool = False, + results_path: Optional[Union[str, Path]] = None, + results_path_extension: Optional[str] = None, + type_time: str = "once", + nb_eval: Optional[int] = None, + nb_samples: Optional[int] = None +) -> Dict: + """ + Simulate a planar PCS model. Save the video and figures. + + Args: + num_segments (int): number of segments of the robot. + type_of_derivation (str, optional): type of derivation ("symbolic" or "numeric"). + Defaults to "symbolic". + type_of_integration (str, optional): scheme of integration ("gauss-legendre", "gauss-kronrad" or "trapezoid"). + Defaults to "gauss-legendre". + param_integration (int, optional): parameter for the integration scheme (number of points for gauss or number of points for trapezoid). + Defaults to 1000 for trapezoid and 5 for gauss. + type_of_jacobian (str, optional): type of jacobian ("explicit" or "autodiff"). + Defaults to "explicit". + strain_selector (Array, optional): selector for the strains as a boolean array. + Defaults to all. + robot_params (Dict[str, Array], optional): parameters of the robot. + strain_selector (Array, optional): selector for the strains. + t (float, optional): time of simulation [s]. + Defaults to 1.0s. + dt (float, optional): time step [s]. + Defaults to 1e-4s. + q0 (Array, optional): initial configuration. + Defaults to a configuration with 5.0pi rad of bending, 0.1 of shear, and 0.2 of axial strain for each segment. + q_d0 (Array, optional): initial velocity. + Defaults to 0 array. + bool_print (bool, optional): if True, print the simulation results. + Defaults to True. + bool_plot (bool, optional): if True, show the figures. + Defaults to True. + bool_save_plot (bool, optional): if True, save the figures. + Defaults to False. + bool_save_video (bool, optional): if True, save the video. + Defaults to False. + bool_save_res (bool, optional): if True, save the results of the simulation. + Defaults to False. + results_path (str, optional): path to save the dictionary with the simulation results. Must have the suffix .pkl. + results_path_extension (str, optional): extension to add to the results path. + Defaults to None, which will use the default path. + type_time (str, optional): type of time evaluation ("once" or "over_time"). + If "once", the time is evaluated only once for the whole simulation. + If "over_time", the time is evaluated for each step of the simulation. + Defaults to "once". + nb_eval (int, optional): number of evaluations for the time evaluation. + Defaults to None, which will use the default number of evaluations (1) + nb_samples (int, optional): number of samples for the time evaluation. + Defaults to None, which will use the default number of samples (1) + + Returns: + Dict: simulation results + - params: parameters of the simulation + - num_segments: number of segments + - type_of_derivation: type of derivation + - type_of_integration: type of integration + - param_integration: parameter for the integration + - robot_params: parameters of the robot + - strain_selector: selector for the strains + - q0: initial configuration + - q_d0: initial velocity + - t: time of simulation + - dt: time step + - video_ts: time steps for video + - results: results of the simulation + - q_ts: configuration trajectory + - q_d_ts: velocity trajectory + - chi_ee_ts: end-effector position trajectory + - U_ts: potential energy trajectory + - T_ts: kinetic energy trajectory + - execution_time: execution time of each step + - import_model: time to import the model + - compile_dynamical_matrices: time to JIT-compile the dynamical matrices function + - evaluate_dynamical_matrices: time to evaluate the dynamical matrices function + - compile_ode: time to JIT-compile the ODE function + - evaluate_ode: time to evaluate the ODE function + - evaluate_forward_kinematics: time to evaluate the forward kinematics function + - evaluate_potential_energy: time to evaluate the potential energy function + - evaluate_kinetic_energy: time to evaluate the kinetic energy function + - draw_robot: time to draw the robot + """ + + # =================================================== + # Initialization of the simulation parameters + # =================================================== + + if not isinstance(num_segments, int): + raise TypeError(f"num_segments must be an integer, but got {type(num_segments).__name__}") + if num_segments < 1: + raise ValueError(f"num_segments must be greater than 0, but got {num_segments}") + + if not isinstance(type_of_derivation, str): + raise TypeError(f"type_of_derivation must be a string, but got {type(type_of_derivation).__name__}") + if type_of_derivation == "numeric": + if not isinstance(type_of_integration, str): + raise TypeError(f"type_of_integration must be a string, but got {type(type_of_integration).__name__}") + if param_integration is None: + if type_of_integration == "gauss-legendre": + param_integration = 5 + if type_of_integration == "gauss-kronrad": + param_integration = 15 + elif type_of_integration == "trapezoid": + param_integration = 1000 + if not isinstance(param_integration, int): + raise TypeError(f"param_integration must be an integer, but got {type(param_integration).__name__}") + if param_integration < 1: + raise ValueError(f"param_integration must be greater than 0, but got {param_integration}") + + if type_of_jacobian not in ["explicit", "autodiff"]: + raise ValueError( + f"type_of_jacobian must be 'explicit' or 'autodiff', but got {type_of_jacobian}" + ) + + elif type_of_derivation == "symbolic": + # filepath to symbolic expressions + sym_exp_filepath = ( + Path(jsrm.__file__).parent + / "symbolic_expressions" + / f"planar_pcs_ns-{num_segments}.dill" + ) + if not sym_exp_filepath.exists(): + raise FileNotFoundError( + f"File {sym_exp_filepath} does not exist. Please run the script to generate the symbolic expressions." + ) + else: + raise ValueError( + f"type_of_derivation must be 'symbolic' or 'numeric', but got {type_of_derivation}" + ) + + if robot_params is None: + # set parameters + rho = 1070 * jnp.ones((num_segments,)) # Volumetric density of Dragon Skin 20 [kg/m^3] + robot_params = { + "th0": jnp.array(0.0), # Initial orientation angle [rad] + "l": 1e-1 * jnp.ones((num_segments,)), # Length of each segment [m] + "r": 2e-2 * jnp.ones((num_segments,)), # Radius of each segment [m] + "rho": rho, + "g": jnp.array([0.0, 9.81]), # Gravity vector [m/s^2] + "E": 2e3 * jnp.ones((num_segments,)), # Elastic modulus [Pa] + "G": 1e3 * jnp.ones((num_segments,)), # Shear modulus [Pa] + } + robot_params["D"] = 1e-3 * jnp.diag( # Damping matrix [Ns/m] + (jnp.repeat( + jnp.array([[1e0, 1e3, 1e3]]), num_segments, axis=0 + ) * robot_params["l"][:, None]).flatten() + ) + #TODO: check if the parameters are correctly defined + + # Max number of degrees of freedom = size of the strain vector + n_xi = 3 * num_segments + + if strain_selector is None: + # activate all strains (i.e. bending, shear, and axial) + strain_selector = jnp.ones((n_xi,), dtype=bool) + if not isinstance(strain_selector, jnp.ndarray): + if isinstance(strain_selector, list): + strain_selector = jnp.array(strain_selector) + else: + raise TypeError(f"strain_selector must be a jnp.ndarray, but got {type(strain_selector).__name__}") + strain_selector = strain_selector.flatten() + if strain_selector.shape[0] != n_xi: + raise ValueError( + f"strain_selector must have the same shape as the strain vector, but got {strain_selector.shape[0]} instead of {n_xi}" + ) + if not jnp.issubdtype(strain_selector.dtype, jnp.bool_): + raise TypeError( + f"strain_selector must be a boolean array, but got {strain_selector.dtype}" + ) + + # number of generalized coordinates + n_dof = jnp.sum(strain_selector) + + if q0 is None: + q0_init = jnp.repeat(jnp.array([5.0 * jnp.pi, 0.1, 0.2])[None, :], num_segments, axis=0).flatten() + q0 = jnp.array([q0_init[i] for i in range(3*num_segments) if strain_selector[i]]) + if not isinstance(q0, jnp.ndarray): + if isinstance(q0, list): + q0 = jnp.array(q0) + else: + raise TypeError(f"q0 must be a jnp.ndarray, but got {type(q0).__name__}") + if q0.shape[0] != n_dof: + raise ValueError( + f"q0 must have the same shape as jnp.sum(strain_selector), but got {q0.shape[0]} instead of {jnp.sum(strain_selector)}" + ) + + if q_d0 is None: + # define initial velocity + q_d0 = jnp.zeros_like(q0) + if not isinstance(q_d0, jnp.ndarray): + if isinstance(q_d0, list): + q_d0 = jnp.array(q_d0) + else: + raise TypeError(f"q_d0 must be a jnp.ndarray, but got {type(q_d0).__name__}") + if q_d0.shape[0] != n_dof: + raise ValueError( + f"q_d0 must have the same shape as q0, but got {q_d0.shape[0]} instead of {n_dof}" + ) + + if not isinstance(t, float): + if isinstance(t, int): + t = float(t) + else: + raise TypeError(f"t must be a float, but got {type(t).__name__}") + if t <= 0: + raise ValueError(f"t must be greater than 0, but got {t}") + + if dt is None: + dt = 1e-4 + if not isinstance(dt, float): + if isinstance(dt, int): + dt = float(dt) + else: + raise TypeError(f"dt must be a float, but got {type(dt).__name__}") + if dt <= 0: + raise ValueError(f"dt must be greater than 0, but got {dt}") + if dt > t: + raise ValueError(f"dt must be less than t, but got {dt} > {t}") + + if not isinstance(bool_save_res, bool): + raise TypeError(f"bool_save_res must be a boolean, but got {type(bool_save_res).__name__}") + + if bool_save_res: + if results_path is None: + results_path_parent = ( + Path(__file__).parent + / "results" + / "times" + / "planar_pcs" + / f"ns-{num_segments}") + file_name = f"{('symb' if type_of_derivation == 'symbolic' else 'num')}" + if type_of_derivation == "numeric": + file_name += f"-{type_of_integration}-{param_integration}-{type_of_jacobian}" + + file_name += f"-{type_time}-{nb_eval}-{nb_samples}" + + if results_path_extension is not None: + if not isinstance(results_path_extension, str): + raise TypeError( + f"results_path_extension must be a string, but got {type(results_path_extension).__name__}" + ) + file_name += f"-{results_path_extension}" + + results_path = (results_path_parent / file_name).with_suffix(".pkl") + + if isinstance(results_path, str) or isinstance(results_path, Path): + results_path_obj = Path(results_path) + if results_path_obj.suffix != ".pkl": + raise ValueError( + f"results_path must have the suffix .pkl, but got {results_path_obj.suffix}" + ) + results_path = results_path_obj + else: + raise TypeError( + f"results_path must be a string, but got {type(results_path).__name__}" + ) + + # create the directory if it does not exist + if not results_path.parent.exists(): + results_path.parent.mkdir(parents=True, exist_ok=True) + + # set simulation parameters + ts = jnp.arange(0.0, t, dt) # time steps + skip_step = 10 # how many time steps to skip in between video frames + video_ts = ts[::skip_step] # time steps for video + + # ===================================================================================================== + # Simulation + # ===================================================================================================== + # save the simulation parameters and results + simulation_dict = { + "params": { + "num_segments": num_segments, + "type_of_derivation": type_of_derivation, + "type_of_integration": type_of_integration, + "param_integration": param_integration, + "type_of_jacobian": type_of_jacobian, + "robot_params": robot_params, + "strain_selector": strain_selector, + "q0": q0, + "q_d0": q_d0, + "t": t, + "dt": dt, + "nb_eval": nb_eval, + "nb_samples": nb_samples, + }, + "execution_time": { + } + } + + print("Number of segments:", num_segments) + print("Type of derivation:", type_of_derivation) + if type_of_derivation == "numeric": + print("Type of integration:", type_of_integration) + print("Parameter for integration:", param_integration) + print("Type of jacobian:", type_of_jacobian) + print() + + # ==================================================== + # Import the planar PCS model depending on the type of derivation + print("Importing the planar PCS model...") + + timer_start = time.time() + if type_of_derivation == "symbolic": + strain_basis, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns = ( + planar_pcs_sym.factory(sym_exp_filepath, strain_selector) + ) + elif type_of_derivation == "numeric": + strain_basis, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns = ( + planar_pcs_num.factory( + num_segments, + strain_selector, + integration_type=type_of_integration, + param_integration=param_integration, + jacobian_type=type_of_jacobian + ) + ) + timer_end = time.time() + + simulation_dict["execution_time"]["import_model"] = timer_end - timer_start + + print(f"Importing the planar PCS model took {timer_end - timer_start:.2e} seconds. \n") + + if type_time == "once": + print("Only compilation and evaluation for one (random) element") + + # ==================================================================================================================================================== + # Dynamical matrices + + dynamical_matrices_fn = partial(dynamical_matrices_fn, robot_params) + + def time_dynamical_matrices_once(): + q = jnp.array(onp.random.randn(*q0.shape)).block_until_ready() + qd = jnp.array(onp.random.randn(*q0.shape)).block_until_ready() + B, C, G, K, D, A = dynamical_matrices_fn(q, qd) + B.block_until_ready() # ensure the matrices are computed + C.block_until_ready() # ensure the matrices are computed + G.block_until_ready() # ensure the matrices are computed + K.block_until_ready() # ensure the matrices are computed + D.block_until_ready() # ensure the matrices are computed + A.block_until_ready() # ensure the matrices are computed + return None + + if EVAL_WITHOUT_JIT: + # No JIT evaluation of the dynamical matrices to capture the time of the evaluation + print("Evaluating the dynamical matrices without JIT ...") + + if nb_eval is None or nb_samples is None or nb_samples == 1: + timer_start = time.time() + time_dynamical_matrices_once() + timer_end = time.time() + + simulation_dict["execution_time"]["evaluate_dynamical_matrices_without_JIT_once"] = timer_end - timer_start + + print(f"Evaluating the dynamical matrices without JIT took {timer_end - timer_start:.2e} seconds. \n") + else: + results_time_dynamical_matrices_once = [timeit.timeit(time_dynamical_matrices_once, number=nb_eval)/nb_eval for _ in range(nb_samples)] + mean_time_dynamical_matrices_once = statistics.mean(results_time_dynamical_matrices_once) + std_time_dynamical_matrices_once = statistics.stdev(results_time_dynamical_matrices_once) + + simulation_dict["execution_time"]["evaluate_dynamical_matrices_without_JIT_once"] = (mean_time_dynamical_matrices_once, std_time_dynamical_matrices_once) + + print(f"Evaluating the dynamical matrices without JIT took {mean_time_dynamical_matrices_once:.2e} seconds +/- {std_time_dynamical_matrices_once:.2e} seconds (mean +/- std) over {nb_samples} samples. \n") + + # JIT the functions + print("JIT-compiling the dynamical matrices function...") + dynamical_matrices_fn = jit(dynamical_matrices_fn) + + # First evaluation of the dynamical matrices to trigger JIT compilation + print(f"Evaluating the dynamical matrices for the first time (JIT-compilation) ...") + timer_start = time.time() + time_dynamical_matrices_once() + timer_end = time.time() + + simulation_dict["execution_time"]["compile_dynamical_matrices_once"] = timer_end - timer_start + + print(f"Evaluating the dynamical matrices for the first time took {timer_end - timer_start:.2e} seconds. \n") + + # Second evaluation of the dynamical matrices to capture the time of the evaluation + print("Evaluating the dynamical matrices for the second time (JIT-evaluation) ...") + + if nb_eval is None or nb_samples is None or nb_samples == 1: + timer_start = time.time() + time_dynamical_matrices_once() + timer_end = time.time() + + simulation_dict["execution_time"]["evaluate_dynamical_matrices_once"] = timer_end - timer_start + + print(f"Evaluating the dynamical matrices for the second time took {timer_end - timer_start:.2e} seconds. \n") + else: + results_time_dynamical_matrices_once = [timeit.timeit(time_dynamical_matrices_once, number=nb_eval)/nb_eval for _ in range(nb_samples)] + mean_time_dynamical_matrices_once = statistics.mean(results_time_dynamical_matrices_once) + std_time_dynamical_matrices_once = statistics.stdev(results_time_dynamical_matrices_once) + + simulation_dict["execution_time"]["evaluate_dynamical_matrices_once"] = (mean_time_dynamical_matrices_once, std_time_dynamical_matrices_once) + + print(f"Evaluating the dynamical matrices for the second time took {mean_time_dynamical_matrices_once:.2e} seconds +/- {std_time_dynamical_matrices_once:.2e} seconds (mean +/- std) over {nb_samples} samples. \n") + + # ==================================================================================================================================================== + # Forward kinematics + + forward_kinematics_fn = partial(forward_kinematics_fn, robot_params) + + def time_forward_kinematics_once(): + q = jnp.array(onp.random.randn(*q0.shape)).block_until_ready() + s = jnp.array(onp.random.randn(1)).block_until_ready() + chi_ee = forward_kinematics_fn(q, s).block_until_ready() + return None + + if EVAL_WITHOUT_JIT: + # No JIT evaluation of the forward kinematics to capture the time of the evaluation + print("Evaluating the forward kinematics without JIT ...") + + if nb_eval is None or nb_samples is None or nb_samples == 1: + timer_start = time.time() + time_forward_kinematics_once() + timer_end = time.time() + + simulation_dict["execution_time"]["evaluate_forward_kinematics_without_JIT__once"] = timer_end - timer_start + + print(f"Evaluating the forward kinematics without JIT took {timer_end - timer_start:.2e} seconds. \n") + else: + results_time_forward_kinematics_once = [timeit.timeit(time_forward_kinematics_once, number=nb_eval)/nb_eval for _ in range(nb_samples)] + mean_time_forward_kinematics_once = statistics.mean(results_time_forward_kinematics_once) + std_time_forward_kinematics_once = statistics.stdev(results_time_forward_kinematics_once) + + simulation_dict["execution_time"]["evaluate_forward_kinematics_without_JIT__once"] = (mean_time_forward_kinematics_once, std_time_forward_kinematics_once) + + print(f"Evaluating the forward kinematics without JIT took {mean_time_forward_kinematics_once:.2e} seconds +/- {std_time_forward_kinematics_once:.2e} seconds (mean +/- std) over {nb_samples} samples. \n") + + # JIT the functions + print("JIT-compiling the forward kinematics function...") + forward_kinematics_fn = jit(forward_kinematics_fn) + + # First evaluation of the forward kinematics to trigger JIT compilation + print("Evaluating the forward kinematics associated with the q configuration (JIT-compilation)...") + + timer_start = time.time() + time_forward_kinematics_once() + timer_end = time.time() + + simulation_dict["execution_time"]["compile_forward_kinematics_once"] = timer_end - timer_start + + print(f"Evaluating the forward kinematics for the first time took {timer_end - timer_start:.2e} seconds. \n") + + # Second evaluation of the forward kinematics to capture the time of the evaluation + print("Evaluating the forward kinematics associated with the q configuration for a second time (JIT-evaluation)...") + + if nb_eval is None or nb_samples is None or nb_samples == 1: + timer_start = time.time() + time_forward_kinematics_once() + timer_end = time.time() + + simulation_dict["execution_time"]["evaluate_forward_kinematics_once"] = timer_end - timer_start + + print(f"Evaluating the forward kinematics for a second time took {timer_end - timer_start:.2e} seconds. \n") + else: + results_time_forward_kinematics_once = [timeit.timeit(time_forward_kinematics_once, number=nb_eval)/nb_eval for _ in range(nb_samples)] + mean_time_forward_kinematics_once = statistics.mean(results_time_forward_kinematics_once) + std_time_forward_kinematics_once = statistics.stdev(results_time_forward_kinematics_once) + + simulation_dict["execution_time"]["evaluate_forward_kinematics_once"] = (mean_time_forward_kinematics_once, std_time_forward_kinematics_once) + + print(f"Evaluating the forward kinematics for a second time took {mean_time_forward_kinematics_once:.2e} seconds +/- {std_time_forward_kinematics_once:.2e} seconds (mean +/- std) over {nb_samples} samples. \n") + + # ==================================================================================================================================================== + # Potential energy + + potential_energy_fn = auxiliary_fns["potential_energy_fn"] + potential_energy_fn = partial(potential_energy_fn, robot_params) + + def time_potential_energy_once(): + q = jnp.array(onp.random.randn(*q0.shape)).block_until_ready() + U = potential_energy_fn(q).block_until_ready() + return None + + if EVAL_WITHOUT_JIT: + # No JIT evaluation of the potential energy to capture the time of the evaluation + print("Evaluating the potential energy without JIT ...") + + if nb_eval is None or nb_samples is None or nb_samples == 1: + timer_start = time.time() + time_potential_energy_once() + timer_end = time.time() + + simulation_dict["execution_time"]["evaluate_potential_energy_without_JIT_once"] = timer_end - timer_start + + print(f"Evaluating the potential energy without JIT took {timer_end - timer_start:.2e} seconds. \n") + else: + results_time_potential_energy_once = [timeit.timeit(time_potential_energy_once, number=nb_eval)/nb_eval for _ in range(nb_samples)] + mean_time_potential_energy_once = statistics.mean(results_time_potential_energy_once) + std_time_potential_energy_once = statistics.stdev(results_time_potential_energy_once) + + simulation_dict["execution_time"]["evaluate_potential_energy_without_JIT_once"] = (mean_time_potential_energy_once, std_time_potential_energy_once) + + print(f"Evaluating the potential energy without JIT took {mean_time_potential_energy_once:.2e} seconds +/- {std_time_potential_energy_once:.2e} seconds (mean +/- std) over {nb_samples} samples. \n") + + # JIT the functions + print("JIT-compiling the potential energy function...") + potential_energy_fn = jit(potential_energy_fn) + + # First evaluation of the potential energy to trigger JIT compilation + print("Evaluating the potential energy associated with the q configuration for the first time (JIT-compilation)...") + + timer_start = time.time() + time_potential_energy_once() + timer_end = time.time() + + simulation_dict["execution_time"]["compile_potential_energy_once"] = timer_end - timer_start + + print(f"Evaluating the potential energy associated with the q configuration for the first time took {timer_end - timer_start:.2e} seconds. \n") + + # Second evaluation of the potential energy to capture the time of the evaluation + print("Evaluating the potential energy associated with the q configuration for a second time (JIT-evaluation)...") + + if nb_eval is None or nb_samples is None or nb_samples == 1: + timer_start = time.time() + time_potential_energy_once() + timer_end = time.time() + + simulation_dict["execution_time"]["evaluate_potential_energy_once"] = timer_end - timer_start + + print(f"Evaluating the potential energy associated with the q configuration for a second time took {timer_end - timer_start:.2e} seconds. \n") + else: + results_time_potential_energy_once = [timeit.timeit(time_potential_energy_once, number=nb_eval)/nb_eval for _ in range(nb_samples)] + mean_time_potential_energy_once = statistics.mean(results_time_potential_energy_once) + std_time_potential_energy_once = statistics.stdev(results_time_potential_energy_once) + + simulation_dict["execution_time"]["evaluate_potential_energy_once"] = (mean_time_potential_energy_once, std_time_potential_energy_once) + + print(f"Evaluating the potential energy associated with the q configuration for a second time took {mean_time_potential_energy_once:.2e} seconds +/- {std_time_potential_energy_once:.2e} seconds (mean +/- std) over {nb_samples} samples. \n") + + # ==================================================================================================================================================== + # Kinetic energy + + kinetic_energy_fn = auxiliary_fns["kinetic_energy_fn"] + kinetic_energy_fn = partial(kinetic_energy_fn, robot_params) + + def time_kinetic_energy_once(): + q = jnp.array(onp.random.randn(*q0.shape)).block_until_ready() + qd = jnp.array(onp.random.randn(*q0.shape)).block_until_ready() + T = kinetic_energy_fn(q, qd).block_until_ready() + return None + + if EVAL_WITHOUT_JIT: + # No JIT evaluation of the kinetic energy to capture the time of the evaluation + print("Evaluating the kinetic energy without JIT ...") + + if nb_eval is None or nb_samples is None or nb_samples == 1: + timer_start = time.time() + time_kinetic_energy_once() + timer_end = time.time() + + simulation_dict["execution_time"]["evaluate_kinetic_energy_without_JIT_once"] = timer_end - timer_start + + print(f"Evaluating the kinetic energy without JIT took {timer_end - timer_start:.2e} seconds. \n") + else: + results_time_kinetic_energy_once = [timeit.timeit(time_kinetic_energy_once, number=nb_eval)/nb_eval for _ in range(nb_samples)] + mean_time_kinetic_energy_once = statistics.mean(results_time_kinetic_energy_once) + std_time_kinetic_energy_once = statistics.stdev(results_time_kinetic_energy_once) + + simulation_dict["execution_time"]["evaluate_kinetic_energy_without_JIT_once"] = (mean_time_kinetic_energy_once, std_time_kinetic_energy_once) + + print(f"Evaluating the kinetic energy without JIT took {mean_time_kinetic_energy_once:.2e} seconds +/- {std_time_kinetic_energy_once:.2e} seconds (mean +/- std) over {nb_samples} samples. \n") + + # JIT the functions + print("JIT-compiling the kinetic energy function...") + kinetic_energy_fn = jit(kinetic_energy_fn) + + # First evaluation of the kinetic energy to trigger JIT compilation + print("Evaluating the kinetic energy associated with the q configuration for the first time (JIT-compilation)...") + + timer_start = time.time() + time_kinetic_energy_once() + timer_end = time.time() + + simulation_dict["execution_time"]["compile_kinetic_energy_once"] = timer_end - timer_start + + print(f"Evaluating the kinetic energy associated with the q configuration took {timer_end - timer_start:.2e} seconds. \n") + + # Second evaluation of the kinetic energy to capture the time of the evaluation + print("Evaluating the kinetic energy associated with the q configuration for a second time (JIT-evaluation)...") + + if nb_eval is None or nb_samples is None or nb_samples == 1: + timer_start = time.time() + time_kinetic_energy_once() + timer_end = time.time() + + simulation_dict["execution_time"]["evaluate_kinetic_energy_once"] = timer_end - timer_start + + print(f"Evaluating the kinetic energy associated with the q configuration for a second time took {timer_end - timer_start:.2e} seconds. \n") + else: + results_time_kinetic_energy_once = [timeit.timeit(time_kinetic_energy_once, number=nb_eval)/nb_eval for _ in range(nb_samples)] + mean_time_kinetic_energy_once = statistics.mean(results_time_kinetic_energy_once) + std_time_kinetic_energy_once = statistics.stdev(results_time_kinetic_energy_once) + + simulation_dict["execution_time"]["evaluate_kinetic_energy_once"] = (mean_time_kinetic_energy_once, std_time_kinetic_energy_once) + + print(f"Evaluating the kinetic energy associated with the q configuration for a second time took {mean_time_kinetic_energy_once:.2e} seconds +/- {std_time_kinetic_energy_once:.2e} seconds (mean +/- std) over {nb_samples} samples. \n") + + elif type_time == "over_time": + print(f"Only in time compilation and evaluation starting from q0 = {q0} and q_d0 = {q_d0}\n") + # ==================================================== + # Parameter for the simulation + x0 = jnp.concatenate([q0, q_d0]) # initial condition + tau = jnp.zeros_like(q0) # torques + + # Create the ODE function + ode_fn = ode_factory(dynamical_matrices_fn, robot_params, tau) + ode_fn = jit(ode_fn) # JIT-compile the ODE function + term = ODETerm(ode_fn) + + # ==================================================================================================================================================== + # Dynamical matrices + + diffeqsolve_fn = partial(diffeqsolve, + term, + solver=Tsit5(), + t0=ts[0], + t1=ts[-1], + dt0=dt, + y0=x0, + max_steps=None, + saveat=SaveAt(ts=video_ts)) + + def time_diffeqsolve_over_time(): + sol = diffeqsolve_fn() + sol.ys.block_until_ready() + return sol + + if EVAL_WITHOUT_JIT: + # No JIT evaluation of the dynamical matrices to capture the time of the evaluation + print("Evaluating the ODE without JIT ...") + + if nb_eval is None or nb_samples is None or nb_samples == 1: + timer_start = time.time() + time_diffeqsolve_over_time() + timer_end = time.time() + + simulation_dict["execution_time"]["evaluate_ode_without_JIT_over_time"] = timer_end - timer_start + + print(f"Evaluating the ODE without JIT took {timer_end - timer_start:.2e} seconds. \n") + else: + results_time_diffeqsolve_over_time = [timeit.timeit(time_diffeqsolve_over_time, number=nb_eval)/nb_eval for _ in range(nb_samples)] + mean_time_diffeqsolve_over_time = statistics.mean(results_time_diffeqsolve_over_time) + std_time_diffeqsolve_over_time = statistics.stdev(results_time_diffeqsolve_over_time) + + simulation_dict["execution_time"]["evaluate_ode_without_JIT_over_time"] = (mean_time_diffeqsolve_over_time, std_time_diffeqsolve_over_time) + + print(f"Evaluating the ODE without JIT took {mean_time_diffeqsolve_over_time:.2e} seconds +/- {std_time_diffeqsolve_over_time:.2e} seconds (mean +/- std) over {nb_samples} samples. \n") + + # JIT the functions + print("JIT-compiling the ODE function...") + diffeqsolve_fn = jit(diffeqsolve_fn) # type: ignore + + # First evaluation of the ODE to trigger JIT compilation + print("Solving the ODE for the first time (JIT-compilation)...") + + timer_start = time.time() + time_diffeqsolve_over_time() + timer_end = time.time() + + simulation_dict["execution_time"]["compile_ode_over_time"] = timer_end - timer_start + + print(f"Solving the ODE for the first time took {timer_end - timer_start:.2e} seconds. \n") + + # Second evaluation of the ODE to capture the time of the evaluation + print("Solving the ODE for the second time (JIT-evaluation)...") + + if nb_eval is None or nb_samples is None or nb_samples == 1: + timer_start = time.time() + time_diffeqsolve_over_time() + timer_end = time.time() + + simulation_dict["execution_time"]["evaluate_ode_over_time"] = timer_end - timer_start + + print(f"Solving the ODE for a second time took {timer_end - timer_start:.2e} seconds. \n") + else: + results_time_diffeqsolve_over_time = [timeit.timeit(time_diffeqsolve_over_time, number=nb_eval)/nb_eval for _ in range(nb_samples)] + mean_time_diffeqsolve_over_time = statistics.mean(results_time_diffeqsolve_over_time) + std_time_diffeqsolve_over_time = statistics.stdev(results_time_diffeqsolve_over_time) + + simulation_dict["execution_time"]["evaluate_ode_over_time"] = (mean_time_diffeqsolve_over_time, std_time_diffeqsolve_over_time) + + print(f"Solving the ODE for a second time took {mean_time_diffeqsolve_over_time:.2e} seconds +/- {std_time_diffeqsolve_over_time:.2e} seconds (mean +/- std) over {nb_samples} samples. \n") + + # ==================================================== + # the evolution of the generalized coordinates and velocities + sol = time_diffeqsolve_over_time() # solve the ODE + q_ts = sol.ys[:, :n_dof].block_until_ready() + q_d_ts = sol.ys[:, n_dof:].block_until_ready() + s_max = jnp.array([jnp.sum(robot_params["l"])]) # maximum length of the robot + + # ==================================================================================================================================================== + # Forward kinematics + + forward_kinematics_fn = partial(forward_kinematics_fn, robot_params) + + def time_forward_kinematics_over_time(): + chi_ee_ts = forward_kinematics_fn( + q_ts, s_max + ).block_until_ready() + return None + + if EVAL_WITHOUT_JIT: + # No JIT evaluation of the forward kinematics to capture the time of the evaluation + print("Evaluating the forward kinematics of the end-effector without JIT ...") + + if nb_eval is None or nb_samples is None or nb_samples == 1: + timer_start = time.time() + time_forward_kinematics_over_time() + timer_end = time.time() + + simulation_dict["execution_time"]["evaluate_forward_kinematics_without_JIT_over_time"] = timer_end - timer_start + + print(f"Evaluating the forward kinematics of the end-effector without JIT took {timer_end - timer_start:.2e} seconds. \n") + + else: + results_time_forward_kinematics_over_time = [timeit.timeit(time_forward_kinematics_over_time, number=nb_eval)/nb_eval for _ in range(nb_samples)] + mean_time_forward_kinematics_over_time = statistics.mean(results_time_forward_kinematics_over_time) + std_time_forward_kinematics_over_time = statistics.stdev(results_time_forward_kinematics_over_time) + + simulation_dict["execution_time"]["evaluate_forward_kinematics_without_JIT_over_time"] = (mean_time_forward_kinematics_over_time, std_time_forward_kinematics_over_time) + + print(f"Evaluating the forward kinematics of the end-effector without JIT took {mean_time_forward_kinematics_over_time:.2e} seconds +/- {std_time_forward_kinematics_over_time:.2e} seconds (mean +/- std) over {nb_samples} samples. \n") + + # JIT the functions + print("JIT-compiling and Vmapping the forward kinematics function...") + forward_kinematics_fn = jit(forward_kinematics_fn) + forward_kinematics_fn = vmap(forward_kinematics_fn, in_axes=(0, None)) + + # First evaluation of the forward kinematics to trigger JIT compilation + print("Evaluating the forward kinematics of the end-effector along the trajectory for the first time (JIT-compilation)...") + + timer_start = time.time() + time_forward_kinematics_over_time() + timer_end = time.time() + + simulation_dict["execution_time"]["compile_forward_kinematics_over_time"] = timer_end - timer_start + + print(f"Evaluating the forward kinematics for the first time took {timer_end - timer_start:.2e} seconds. \n") + + # Second evaluation of the forward kinematics to capture the time of the evaluation + print("Evaluating the forward kinematics of the end-effector along the trajectory for a second time (JIT-evaluation)...") + + if nb_eval is None or nb_samples is None or nb_samples == 1: + timer_start = time.time() + time_forward_kinematics_over_time() + timer_end = time.time() + + simulation_dict["execution_time"]["evaluate_forward_kinematics_over_time"] = timer_end - timer_start + + print(f"Evaluating the forward kinematics for a second time took {timer_end - timer_start:.2e} seconds. \n") + else: + results_time_forward_kinematics_over_time = [timeit.timeit(time_forward_kinematics_over_time, number=nb_eval)/nb_eval for _ in range(nb_samples)] + mean_time_forward_kinematics_over_time = statistics.mean(results_time_forward_kinematics_over_time) + std_time_forward_kinematics_over_time = statistics.stdev(results_time_forward_kinematics_over_time) + + simulation_dict["execution_time"]["evaluate_forward_kinematics_over_time"] = (mean_time_forward_kinematics_over_time, std_time_forward_kinematics_over_time) + + print(f"Evaluating the forward kinematics for a second time took {mean_time_forward_kinematics_over_time:.2e} seconds +/- {std_time_forward_kinematics_over_time:.2e} seconds (mean +/- std) over {nb_samples} samples. \n") + + + # ==================================================================================================================================================== + # Potential energy + + potential_energy_fn = auxiliary_fns["potential_energy_fn"] + potential_energy_fn = partial(potential_energy_fn, robot_params) + + def time_potential_energy_over_time(): + U_ts = potential_energy_fn(q_ts).block_until_ready() + return None + + if EVAL_WITHOUT_JIT: + # No JIT evaluation of the potential energy to capture the time of the evaluation + print("Evaluating the potential energy without JIT ...") + + if nb_eval is None or nb_samples is None or nb_samples == 1: + timer_start = time.time() + time_potential_energy_over_time() + timer_end = time.time() + + simulation_dict["execution_time"]["evaluate_potential_energy_without_JIT_over_time"] = timer_end - timer_start + + print(f"Evaluating the potential energy without JIT took {timer_end - timer_start:.2e} seconds. \n") + else: + results_time_potential_energy_over_time = [timeit.timeit(time_potential_energy_over_time, number=nb_eval)/nb_eval for _ in range(nb_samples)] + mean_time_potential_energy_over_time = statistics.mean(results_time_potential_energy_over_time) + std_time_potential_energy_over_time = statistics.stdev(results_time_potential_energy_over_time) + + simulation_dict["execution_time"]["evaluate_potential_energy_without_JIT_over_time"] = (mean_time_potential_energy_over_time, std_time_potential_energy_over_time) + + print(f"Evaluating the potential energy without JIT took {mean_time_potential_energy_over_time:.2e} seconds +/- {std_time_potential_energy_over_time:.2e} seconds (mean +/- std) over {nb_samples} samples. \n") + + # JIT the functions + print("JIT-compiling and Vmapping the potential energy function...") + potential_energy_fn = jit(potential_energy_fn) + potential_energy_fn = vmap(potential_energy_fn) + + # First evaluation of the potential energy to trigger JIT compilation + print("Evaluating the potential energy for the first time (JIT-compilation)...") + + timer_start = time.time() + time_potential_energy_over_time() + timer_end = time.time() + + simulation_dict["execution_time"]["compile_potential_energy_over_time"] = timer_end - timer_start + + print(f"Evaluating the potential energy took {timer_end - timer_start:.2e} seconds. \n") + + # Second evaluation of the potential energy to capture the time of the evaluation + print("Evaluating the potential energy for a second time (JIT-evaluation)...") + + if nb_eval is None or nb_samples is None or nb_samples == 1: + timer_start = time.time() + time_potential_energy_over_time() + timer_end = time.time() + + simulation_dict["execution_time"]["evaluate_potential_energy_over_time"] = timer_end - timer_start + + print(f"Evaluating the potential energy for a second time took {timer_end - timer_start:.2e} seconds. \n") + else: + results_time_potential_energy_over_time = [timeit.timeit(time_potential_energy_over_time, number=nb_eval)/nb_eval for _ in range(nb_samples)] + mean_time_potential_energy_over_time = statistics.mean(results_time_potential_energy_over_time) + std_time_potential_energy_over_time = statistics.stdev(results_time_potential_energy_over_time) + + simulation_dict["execution_time"]["evaluate_potential_energy_over_time"] = (mean_time_potential_energy_over_time, std_time_potential_energy_over_time) + + print(f"Evaluating the potential energy for a second time took {mean_time_potential_energy_over_time:.2e} seconds +/- {std_time_potential_energy_over_time:.2e} seconds (mean +/- std) over {nb_samples} samples. \n") + + # ==================================================================================================================================================== + # Kinetic energy + + kinetic_energy_fn = auxiliary_fns["kinetic_energy_fn"] + kinetic_energy_fn = partial(kinetic_energy_fn, robot_params) + + def time_kinetic_energy_over_time(): + T_ts = kinetic_energy_fn(q_ts, q_d_ts).block_until_ready() + return None + + if EVAL_WITHOUT_JIT: + # No JIT evaluation of the kinetic energy to capture the time of the evaluation + print("Evaluating the kinetic energy without JIT ...") + + if nb_eval is None or nb_samples is None or nb_samples == 1: + timer_start = time.time() + time_kinetic_energy_over_time() + timer_end = time.time() + + simulation_dict["execution_time"]["evaluate_kinetic_energy_without_JIT_over_time"] = timer_end - timer_start + + print(f"Evaluating the kinetic energy without JIT took {timer_end - timer_start:.2e} seconds. \n") + else: + results_time_kinetic_energy_over_time = [timeit.timeit(time_kinetic_energy_over_time, number=nb_eval)/nb_eval for _ in range(nb_samples)] + mean_time_kinetic_energy_over_time = statistics.mean(results_time_kinetic_energy_over_time) + std_time_kinetic_energy_over_time = statistics.stdev(results_time_kinetic_energy_over_time) + + simulation_dict["execution_time"]["evaluate_kinetic_energy_without_JIT_over_time"] = (mean_time_kinetic_energy_over_time, std_time_kinetic_energy_over_time) + + print(f"Evaluating the kinetic energy without JIT took {mean_time_kinetic_energy_over_time:.2e} seconds +/- {std_time_kinetic_energy_over_time:.2e} seconds (mean +/- std) over {nb_samples} samples. \n") + + # JIT the functions + print("JIT-compiling and Vmapping the kinetic energy function...") + kinetic_energy_fn = jit(kinetic_energy_fn) + kinetic_energy_fn = vmap(kinetic_energy_fn) + + # First evaluation of the kinetic energy to trigger JIT compilation + print("Evaluating the kinetic energy for the first time (JIT-compilation)...") + + timer_start = time.time() + time_kinetic_energy_over_time() + timer_end = time.time() + + simulation_dict["execution_time"]["compile_kinetic_energy_over_time"] = timer_end - timer_start + + print(f"Evaluating the kinetic energy took {timer_end - timer_start:.2e} seconds. \n") + + # Second evaluation of the kinetic energy to capture the time of the evaluation + print("Evaluating the kinetic energy for a second time (JIT-evaluation)...") + + if nb_eval is None or nb_samples is None or nb_samples == 1: + timer_start = time.time() + time_kinetic_energy_over_time() + timer_end = time.time() + + simulation_dict["execution_time"]["evaluate_kinetic_energy_over_time"] = timer_end - timer_start + + print(f"Evaluating the kinetic energy for a second time took {timer_end - timer_start:.2e} seconds. \n") + else: + results_time_kinetic_energy_over_time = [timeit.timeit(time_kinetic_energy_over_time, number=nb_eval)/nb_eval for _ in range(nb_samples)] + mean_time_kinetic_energy_over_time = statistics.mean(results_time_kinetic_energy_over_time) + std_time_kinetic_energy_over_time = statistics.stdev(results_time_kinetic_energy_over_time) + + simulation_dict["execution_time"]["evaluate_kinetic_energy_over_time"] = (mean_time_kinetic_energy_over_time, std_time_kinetic_energy_over_time) + + print(f"Evaluating the kinetic energy for a second time took {mean_time_kinetic_energy_over_time:.2e} seconds +/- {std_time_kinetic_energy_over_time:.2e} seconds (mean +/- std) over {nb_samples} samples. \n") + + # =========================================================================== + # Save the simulation time results + # =========================== + if bool_save_res: + print("Saving the simulation results...") + assert results_path is not None, "results_path should not be None when saving results" + with open(results_path, "wb") as f: + pickle.dump(simulation_dict, f) + print(f"Simulation results saved at {results_path} \n") + else: + print("Simulation results not saved. \n") + + print("Simulation finished. \n") + print("**************************************************************** \n \n") + + return simulation_dict + +if __name__ == "__main__": + simulate_planar_pcs_value_eval(num_segments=1) diff --git a/examples/simulate_planar_pcs.py b/examples/simulate_planar_pcs.py index 92e962e..e71eeac 100644 --- a/examples/simulate_planar_pcs.py +++ b/examples/simulate_planar_pcs.py @@ -126,6 +126,8 @@ def draw_robot( tau = jnp.zeros_like(q0) # torques ode_fn = ode_factory(dynamical_matrices_fn, params, tau) + # jit the ODE function + ode_fn = jax.jit(ode_fn) term = ODETerm(ode_fn) sol = diffeqsolve( @@ -145,10 +147,14 @@ def draw_robot( # the evolution of the generalized velocities q_d_ts = sol.ys[:, n_q:] + s_max = jnp.array([jnp.sum(params["l"])]) + + forward_kinematics_fn_end_effector = partial(forward_kinematics_fn, params, s=s_max) + forward_kinematics_fn_end_effector = jax.jit(forward_kinematics_fn_end_effector) + forward_kinematics_fn_end_effector = vmap(forward_kinematics_fn_end_effector) + # evaluate the forward kinematics along the trajectory - chi_ee_ts = vmap(forward_kinematics_fn, in_axes=(None, 0, None))( - params, q_ts, jnp.array([jnp.sum(params["l"])]) - ) + chi_ee_ts = forward_kinematics_fn_end_effector(q_ts) # plot the configuration vs time plt.figure() for segment_idx in range(num_segments): @@ -202,10 +208,10 @@ def draw_robot( # plot the energy along the trajectory kinetic_energy_fn_vmapped = vmap( - partial(auxiliary_fns["kinetic_energy_fn"], params) + partial(jax.jit(auxiliary_fns["kinetic_energy_fn"]), params) ) potential_energy_fn_vmapped = vmap( - partial(auxiliary_fns["potential_energy_fn"], params) + partial(jax.jit(auxiliary_fns["potential_energy_fn"]), params) ) U_ts = potential_energy_fn_vmapped(q_ts) T_ts = kinetic_energy_fn_vmapped(q_ts, q_d_ts) diff --git a/examples/videos/.gitignore b/examples/videos/.gitignore deleted file mode 100644 index 2641667..0000000 --- a/examples/videos/.gitignore +++ /dev/null @@ -1 +0,0 @@ -*.mp4 diff --git a/pyproject.toml b/pyproject.toml index 99c08f1..c13dec0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ name = "jsrm" # Required # # For a discussion on single-sourcing the version, see # https://packaging.python.org/guides/single-sourcing-package-version/ -version = "0.0.16" # Required +version = "0.0.17" # Required # This is a one-line description or tagline of what your project does. This # corresponds to the "Summary" metadata field: @@ -57,7 +57,8 @@ keywords = ["JAX", "Soft Robotics", "Kinematics", "Dynamics"] # Optional # authored the project, and a valid email address corresponding to the name # listed. authors = [ - {name = "Maximilian Stölzle", email = "maximilian@stoelzle.ch" } # Optional + {name = "Maximilian Stölzle", email = "maximilian@stoelzle.ch" }, + {name = "Solange Gribonval", email = "solange.gribonval@polytechnique.edu" }, ] # This should be your name or the names of the organization who currently @@ -105,6 +106,7 @@ dependencies = [ # Optional "dill", "jax", "numpy", + "quadax", "peppercorn", "sympy>=1.11" ] diff --git a/src/jsrm/integration.py b/src/jsrm/integration.py index e33bd78..2a3a0f6 100644 --- a/src/jsrm/integration.py +++ b/src/jsrm/integration.py @@ -1,4 +1,4 @@ -from jax import Array, jit +from jax import Array from typing import Callable, Dict from jsrm.systems import euler_lagrangian @@ -26,7 +26,6 @@ def ode_factory( ode_fn: ODE function of the form ode_fn(t, x) -> x_dot """ - @jit def ode_fn(t: float, x: Array, *args) -> Array: """ ODE of the dynamical Lagrangian system. @@ -69,7 +68,6 @@ def ode_with_forcing_factory( ode_fn: ODE function of the form ode_fn(t, x, tau) -> x_dot """ - @jit def ode_fn( t: float, x: Array, diff --git a/src/jsrm/math_utils.py b/src/jsrm/math_utils.py index 3c21955..8088921 100644 --- a/src/jsrm/math_utils.py +++ b/src/jsrm/math_utils.py @@ -1,11 +1,12 @@ from jax import numpy as jnp -from jax import Array, lax, jit +from jax import Array, lax - -@jit -def blk_diag(a: Array) -> Array: +def blk_diag( + a: Array +) -> Array: """ - Create a block diagonal matrix from a tensor of blocks + Create a block diagonal matrix from a tensor of blocks. + Args: a: matrices to be block diagonalized of shape (m, n, o) @@ -31,7 +32,7 @@ def assign_block_diagonal(i, _b): # Implement for loop to assign each block in `a` to the block-diagonal of `b` # Hint: use `jax.lax.fori_loop` and pass `assign_block_diagonal` as an argument - b = jnp.zeros((a.shape[0] * a.shape[1], a.shape[0] * a.shape[2])) + b = jnp.zeros((a.shape[0] * a.shape[1], a.shape[0] * a.shape[2]), dtype=a.dtype) b = lax.fori_loop( lower=0, upper=a.shape[0], @@ -40,3 +41,32 @@ def assign_block_diagonal(i, _b): ) return b + +def blk_concat( + a: Array +) -> Array: + """ + Concatenate horizontally (along the columns) a list of N matrices of size (m, n) to create a single matrix of size (m, n * N). + + Args: + a (Array): matrices to be concatenated of shape (N, m, n) + + Returns: + b (Array): concatenated matrix of shape (m, N * n) + """ + b = a.transpose(1, 0, 2).reshape(a.shape[1], -1) + return b + +if __name__ == "__main__": + # Example usage + a = jnp.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) + print("Original array:") + print(a) + + b = blk_diag(a) + print("Block diagonal matrix:") + print(b) + + c = blk_concat(a) + print("Concatenated matrix:") + print(c) \ No newline at end of file diff --git a/src/jsrm/symbolic_derivation/planar_pcs.py b/src/jsrm/symbolic_derivation/planar_pcs.py index 93d7c19..3973cf3 100644 --- a/src/jsrm/symbolic_derivation/planar_pcs.py +++ b/src/jsrm/symbolic_derivation/planar_pcs.py @@ -65,9 +65,9 @@ def symbolically_derive_planar_pcs_model( # tendon length jacobians for each segment as a function of the point coordinate s J_tend_sms = [] # cross-sectional area of each segment - A = sp.zeros(num_segments) + A = sp.zeros(num_segments, 1) # second area moment of inertia of each segment - I = sp.zeros(num_segments) + I = sp.zeros(num_segments, 1) # inertia matrix B = sp.zeros(num_dof, num_dof) # potential energy diff --git a/src/jsrm/systems/__init__.py b/src/jsrm/systems/__init__.py index e69de29..ede258c 100644 --- a/src/jsrm/systems/__init__.py +++ b/src/jsrm/systems/__init__.py @@ -0,0 +1 @@ +import jsrm.systems.planar_pcs_sym as planar_pcs diff --git a/src/jsrm/systems/planar_pcs_num.py b/src/jsrm/systems/planar_pcs_num.py new file mode 100644 index 0000000..8021a18 --- /dev/null +++ b/src/jsrm/systems/planar_pcs_num.py @@ -0,0 +1,1768 @@ +from jax import Array, lax, vmap +from jax import jacobian, grad +from jax import scipy as jscipy +from jax import numpy as jnp +from quadax import GaussKronrodRule + +import numpy as onp +from typing import Callable, Dict, Tuple, Optional, Literal, Any + +from .utils import ( + compute_strain_basis, + compute_planar_stiffness_matrix, + gauss_quadrature, +) +from jsrm.math_utils import blk_diag, blk_concat +from jsrm.utils.lie_operators import ( # To use SE(2) + Tangent_gn_SE2, + Adjoint_gn_SE2_inv, + Adjoint_g_SE2, + adjoint_SE2, + adjoint_star_SE2, +) +from jsrm.utils.lie_operators import ( + compute_weighted_sums, +) + +# To extract the interest coordinates and strains from the SE(3) elements +INTEREST_COORDINATES = jnp.array( + [2, 3, 4] +) # indices of the interest coordinates in the SE3 forward kinematics vector [theta_x, theta_y, theta_z, x, y, z] => [theta_z, x, y] +INTEREST_STRAIN = jnp.array( + [2, 3, 4] +) # indices of the interest strains in the SE3 strain vector [kappa_x, kappa_y, kappa_z, sigma_x, sigma_y, sigma_z] => [kappa_z, sigma_x, sigma_y] +# To reorder the lines to match the forward kinematics vector or strain vector +REORDERED_LINES_FWD_KINE = jnp.array( + [1, 2, 0] +) # reorder the lines to match the forward kinematics vector [x, y, theta] => [sigma_x, sigma_y, kappa_z] +REORDERED_LINES_STRAIN = jnp.array( + [2, 0, 1] +) # reorder the lines to match the strain vector [kappa_z, sigma_x, sigma_y] => [theta, x, y] + + +def factory( + num_segments: int, + strain_selector: Optional[Array] = None, + xi_eq: Optional[Array] = None, + stiffness_fn: Optional[Callable] = None, + actuation_mapping_fn: Optional[Callable] = None, + global_eps: float = jnp.finfo(jnp.float32).eps, + integration_type: Optional[ + Literal["gauss-legendre", "gauss-kronrad", "trapezoid"] + ] = "gauss-legendre", + param_integration: Optional[int] = None, + jacobian_type: Optional[Literal["explicit", "autodiff"]] = "explicit", +) -> Tuple[ + Array, + Callable[[Dict[str, Array], Array, Array, float], Array], + Callable[ + [Dict[str, Array], Array, Array, float], + Tuple[Array, Array, Array, Array, Array, Array], + ], + Dict[str, Callable[..., Any]], +]: + """ + Factory function to create the forward kinematics function for a planar robot. + This function computes the forward kinematics of a planar robot with a given number of segments. + + Args: + num_segments (int): number of segments in the robot. + strain_selector (Array, optional): strain selector array of shape (3 * num_segments, ) + specifying which strain components are active by setting them to True or False. + Defaults to None. + xi_eq (Array, optional): equilibrium strain vector of shape (3 * num_segments, ). + Defaults to 1 for the axial strain and 0 for the bending and shear strains. + stiffness_fn (Callable, optional): function to compute the stiffness matrix. + Defaults to None. + actuation_mapping_fn (Callable, optional): function to compute the actuation mapping. + Defaults to None. + global_eps (float, optional): small number to avoid singularities. + Defaults to 1e-8. + integration_type (str, optional): type of integration to use: "gauss-legendre", "gauss-kronrad" or "trapezoid". + Defaults to "gauss-legendre" for Gaussian quadrature. + param_integration (int, optional): parameter for the integration method. + If None, it is set to 5 for Gauss-Legendre quadrature, 15 for Gauss-Kronrad quadrature and 1000 for trapezoidal integration. + jacobian_type (str, optional): type of Jacobian to compute: "explicit" or "autodiff". + Defaults to "explicit" for explicit Jacobian computation. + + Returns: + B_xi (Array): strain basis matrix of shape (n_xi, n_q) where n_xi is the number of strains and n_q is the number of configuration variables. + forward_kinematics_fn (Callable): function to compute the forward kinematics of the robot. + takes in robot parameters params, configuration vector q, and point coordinate s along the robot, + and returns the pose of the robot at a given point along its length. + dynamical_matrices_fn (Callable): function to compute the dynamical matrices of the robot. + takes in robot parameters params, configuration vector q, configuration velocity q_d, + and returns the dynamical matrices B, C, G, K, D, alpha + auxiliary_fns (Dict[str, Callable]): dictionary of auxiliary functions for the robot. + - "apply_eps_to_bend_strains": function to apply a small number to the bending strains to avoid singularities. + - "classify_segment": function to classify a point along the robot to the corresponding segment. + - "stiffness_fn": function to compute the stiffness matrix of the robot. + - "actuation_mapping_fn": actuation_mapping_fn, + - "jacobian_fn": inertial-frame Jacobian of the forward kinematics function with respect to the strain vector. + - "kinetic_energy_fn": kinetic energy function of the robot. + - "potential_energy_fn": potential energy function of the robot. + - "energy_fn": total energy function of the robot. + - "operational_space_dynamical_matrices_fn": function to compute the operational space dynamical matrices of the robot. + """ + + # ======================================================================================================================= + # Initialize parameters if not provided + # ==================================================== + # Number of segments + if not isinstance(num_segments, int): + raise ValueError( + f"num_segments must be an integer, but got {type(num_segments)}" + ) + if num_segments < 1: + raise ValueError(f"num_segments must be greater than 0, but got {num_segments}") + + # Max number of degrees of freedom = size of the strain vector + n_xi = 3 * num_segments + + # Strain basis matrix + if strain_selector is None: + # activate all strains (i.e. bending, shear, and axial) + strain_selector = jnp.ones((n_xi,), dtype=bool) + if not isinstance(strain_selector, jnp.ndarray): + if isinstance(strain_selector, list): + strain_selector = jnp.array(strain_selector) + else: + raise TypeError( + f"strain_selector must be a jnp.ndarray, but got {type(strain_selector).__name__}" + ) + strain_selector = strain_selector.flatten() + if strain_selector.shape[0] != n_xi: + raise ValueError( + f"strain_selector must have the same shape as the strain vector, but got {strain_selector.shape[0]} instead of {n_xi}" + ) + if not jnp.issubdtype(strain_selector.dtype, jnp.bool_): + raise TypeError( + f"strain_selector must be a boolean array, but got {strain_selector.dtype}" + ) + + # Rest strain + if xi_eq is None: + xi_eq = jnp.zeros((n_xi,)) + # By default, set the axial rest strain (local y-axis) along the entire rod to 1.0 + rest_strain_reshaped = xi_eq.reshape((-1, 3)) + rest_strain_reshaped = rest_strain_reshaped.at[:, -1].set(1.0) + xi_eq = rest_strain_reshaped.flatten() + if not isinstance(xi_eq, jnp.ndarray): + if isinstance(xi_eq, list): + xi_eq = jnp.array(xi_eq) + else: + raise TypeError( + f"xi_eq must be a jnp.ndarray, but got {type(xi_eq).__name__}" + ) + xi_eq = xi_eq.flatten() + if xi_eq.shape[0] != n_xi: + raise ValueError( + f"xi_eq must have the same shape as the strain vector, but got {xi_eq.shape[0]} instead of {n_xi}" + ) + if not jnp.issubdtype(xi_eq.dtype, jnp.floating): + if not jnp.issubdtype(xi_eq.dtype, jnp.integer): + raise TypeError( + f"xi_eq must be a floating point array, but got {xi_eq.dtype}" + ) + else: + xi_eq = xi_eq.astype(jnp.float32) + + # Stiffness function + compute_stiffness_matrix_for_all_segments_fn = vmap(compute_planar_stiffness_matrix) + if stiffness_fn is None: + + def stiffness_fn( + params: Dict[str, Array], + B_xi: Array, + formulate_in_strain_space: bool = False, + ) -> Array: + """ + Compute the stiffness matrix of the system. + Args: + params: Dictionary of robot parameters + B_xi: Strain basis matrix + formulate_in_strain_space: + whether to formulate the elastic matrix in the strain space + Returns: + S: elastic matrix of shape (n_q, n_q) if formulate_in_strain_space is False or (n_xi, n_xi) otherwise + """ + # length of the segments + l = params["l"] + # cross-sectional area and second moment of area + A = jnp.pi * params["r"] ** 2 + Ib = A**2 / (4 * jnp.pi) + + # elastic and shear modulus + E, G = params["E"], params["G"] + # stiffness matrix of shape (num_segments, 3, 3) + S_sms = compute_stiffness_matrix_for_all_segments_fn(l, A, Ib, E, G) + # we define the elastic matrix of shape (n_xi, n_xi) as K(xi) = K @ xi where K is equal to + S = blk_diag(S_sms) + + if not formulate_in_strain_space: + S = B_xi.T @ S @ B_xi + + return S + + if not callable(stiffness_fn): + raise TypeError( + f"stiffness_fn must be a callable, but got {type(stiffness_fn).__name__}" + ) + + # Actuation mapping function + if actuation_mapping_fn is None: + + def actuation_mapping_fn( + forward_kinematics_fn: Callable, + jacobian_fn: Callable, + params: Dict[str, Array], + B_xi: Array, + q: Array, + eps: float = global_eps, + ) -> Array: + """ + Returns the actuation matrix that maps the actuation space to the configuration space. + Assumes the fully actuated and identity actuation matrix. + Args: + forward_kinematics_fn: function to compute the forward kinematics + jacobian_fn: function to compute the Jacobian + params: dictionary with robot parameters + B_xi: strain basis matrix + q: configuration of the robot + eps: small number to avoid singularities (default: global_eps = 1e-8) + Returns: + A: actuation matrix of shape (n_xi, n_xi) where n_xi is the number of strains. + """ + A = B_xi.T @ jnp.identity(n_xi) @ B_xi + + return A + + if not callable(actuation_mapping_fn): + raise TypeError( + f"actuation_mapping_fn must be a callable, but got {type(actuation_mapping_fn).__name__}" + ) + + if integration_type == "gauss-legendre": + if param_integration is None: + param_integration = 5 + elif integration_type == "gauss-kronrad": + if param_integration is None: + param_integration = 15 + if param_integration not in [15, 21, 31, 41, 51, 61]: + raise ValueError( + f"param_integration must be one of [15, 21, 31, 41, 51, 61] for gauss-kronrad integration, but got {param_integration}" + ) + elif integration_type == "trapezoid": + if param_integration is None: + param_integration = 1000 + else: + raise ValueError( + f"integration_type must be either 'gauss-legendre', 'gauss-kronrad' or 'trapezoid', but got {integration_type}" + ) + + if jacobian_type not in ["explicit", "autodiff"]: + raise ValueError( + f"jacobian_type must be either 'explicit' or 'autodiff', but got {jacobian_type}" + ) + + # ======================================================================================================================= + # Define the functions + # ==================================================== + + # Compute the strain basis matrix + B_xi = compute_strain_basis(strain_selector) + + def apply_eps_to_bend_strains(xi: Array, eps: float) -> Array: + """ + Add a small number to the bending strain to avoid singularities. + + Args: + xi (Array): strain vector of the robot. + eps (float): small number to add to the bending strain. + + Returns: + Array: strain vector with the bending strain modified. + """ + if eps == None: + return xi + else: + xi_reshaped = xi.reshape((-1, 3)) + + xi_bend_sign = jnp.sign(xi_reshaped[:, 0]) + + # set zero sign to 1 (i.e. positive) + xi_bend_sign = jnp.where(xi_bend_sign == 0, 1, xi_bend_sign) + + # add eps to the bending strain (i.e. the first column) + sigma_b_epsed = lax.select( + jnp.abs(xi_reshaped[:, 0]) < eps, + xi_bend_sign * eps, + xi_reshaped[:, 0], + ) + xi_epsed = jnp.stack( + [ + sigma_b_epsed, + xi_reshaped[:, 1], + xi_reshaped[:, 2], + ], + axis=1, + ) + + # Flatten the array + xi_epsed = xi_epsed.flatten() + + return xi_epsed + + def classify_segment( + params: Dict[str, Array], s: Array + ) -> Tuple[Array, Array, Array]: + """ + Classify the point along the robot to the corresponding segment. + + Args: + params (Dict[str, Array]): dictionary of robot parameters + s (Array): point coordinate along the robot in the interval [0, L]. + + Returns: + segment_idx (Array): index of the segment where the point is located + s_segment (Array): point coordinate along the segment in the interval [0, l_segment] + l_cum (Array): cumulative length of the segments starting with 0 + """ + l = params["l"] + + # Compute the cumulative length of the segments starting with 0 + l_cum = jnp.cumsum(jnp.concatenate([jnp.zeros(1), l])) + + # Classify the point along the robot to the corresponding segment + segment_idx = jnp.clip(jnp.sum(s > l_cum) - 1, 0, len(l) - 1) + + # Compute the point coordinate along the segment in the interval [0, l_segment] + s_segment = s - l_cum[segment_idx] + + return segment_idx, s_segment.squeeze(), l_cum + + def chi_fn( + params: Dict[str, Array], xi: Array, s: Array, eps: float = global_eps + ) -> Array: + """ + Compute the pose of the robot at a given point along its length with respect to the strain vector. + + Args: + params (Dict[str, Array]): dictionary of robot parameters. + xi (Array): strain vector of the robot. + s (Array): point coordinate along the robot in the interval [0, L]. + eps (float, optional): small number to avoid singularities. Defaults to global_eps = 1e-8. + + Returns: + chi (Array): pose of the robot at the point s in the interval [0, L] + The pose is represented as a vector [x, y, theta] where (x, y) is the position + and theta is the orientation angle. + """ + th0 = params["th0"] # initial angle of the robot + l = params["l"] # length of each segment [m] + + # Classify the point along the robot to the corresponding segment + segment_idx, s_local, _ = classify_segment(params, s) + + chi_0 = jnp.concatenate([jnp.zeros(2), th0[None]]) # Initial pose of the robot + + # Iteration function + def chi_i(i: int, chi_prev: Array) -> Array: + th_prev = chi_prev[2] # Extract previous orientation angle from val + p_prev = chi_prev[:2] # Extract previous position from val + + # Extract strains for the current segment + kappa = xi[3 * i + 0] # Bending strain + sigma_x = xi[3 * i + 1] # Shear strain + sigma_y = xi[3 * i + 2] # Axial strain + + # Compute the length of the current segment to integrate + l_i = jnp.where(i == segment_idx, s_local, l[i]) + + # Compute the orientation angle for the current segment + dth = kappa * l_i # Angle increment for the current segment + th = th_prev + dth + + # Compute the integrals for the transformation matrix + int_cos_th = jnp.where( + jnp.abs(kappa) < eps, + l_i * jnp.cos(th_prev), + (jnp.sin(th) - jnp.sin(th_prev)) / kappa, + ) + + int_sin_th = jnp.where( + jnp.abs(kappa) < eps, + l_i * jnp.sin(th_prev), + (jnp.cos(th_prev) - jnp.cos(th)) / kappa, + ) + + # Transformation matrix + R = jnp.stack( + [ + jnp.stack([int_cos_th, -int_sin_th]), + jnp.stack([int_sin_th, int_cos_th]), + ] + ) + + # Compute the position + p = p_prev + R @ jnp.stack([sigma_x, sigma_y], axis=-1) + + return jnp.concatenate([p, th[None]]) + + _, chi_list = lax.scan( + f=lambda carry, i: (chi_i(i, carry), chi_i(i, carry)), + init=chi_0, + xs=jnp.arange(num_segments + 1), + ) + + chi = chi_list[segment_idx] + + return chi + + def forward_kinematics_fn( + params: Dict[str, Array], q: Array, s: Array, eps: float = global_eps + ) -> Array: + """ + Compute the forward kinematics of the robot. + + Args: + params (Dict[str, Array]): dictionary of robot parameters. + q (Array): configuration vector of the robot. + s (Array): point coordinate along the robot in the interval [0, L]. + eps (float, optional): small number to avoid singularities. Defaults to global_eps = 1e-8. + + Returns: + chi (Array): pose of the robot at the point s in the interval [0, L]. + The pose is represented as a vector [x, y, theta] where (x, y) is the position + and theta is the orientation angle. + """ + # Map the configuration to the strains + xi = xi_eq + B_xi @ q + + chi = chi_fn(params, xi, s, eps) + + return chi + + def J_autodiff( + params: Dict[str, Array], xi: Array, s: Array, eps: float = global_eps + ) -> Array: + """ + Compute the inertial-frame jacobian of the forward kinematics function with respect to the strain vector + using autodiff. + + Args: + params (Dict[str, Array]): dictionary of robot parameters. + xi (Array): strain vector of the robot. + s (Array): point coordinate along the robot in the interval [0, L]. + eps (float, optional): small number to avoid singularities. Defaults to global_eps = 1e-8. + + Returns: + J (Array): inertial-frame jacobian of the forward kinematics function with respect to the strain vector. + """ + # Compute the Jacobian of chi_fn with respect to xi + J = jacobian(lambda _xi: chi_fn(params, _xi, s, eps))(xi) + + # apply the strain basis to the Jacobian + J = J @ B_xi + + return J + + def J_explicit_local( + params: Dict[str, Array], xi: Array, s: Array, eps: float = global_eps + ) -> Array: + """ + Compute the body-frame jacobian of the forward kinematics function with respect to the strain vector + at a given point s using explicit expression in SE(2). + + Args: + params (Dict[str, Array]): dictionary of robot parameters. + xi (Array): strain vector of the robot. + s (Array): point coordinate along the robot in the interval [0, L_tot]. + + Returns: + J_local (Array) : body-frame jacobian of the forward kinematics function with respect to the strain vector + using explicit expression in SE(2). + """ + + # Classify the point along the robot to the corresponding segment + segment_idx, _, l_cum = classify_segment(params, s) + + xi = apply_eps_to_bend_strains(xi, eps) # Apply eps to the bending strain + + # Initial condition + xi_SE2_0 = xi[0:3] + + Ad_g0_inv_L0 = Adjoint_gn_SE2_inv(xi_SE2_0, l_cum[0], l_cum[1]) + Ad_g0_inv_s = Adjoint_gn_SE2_inv(xi_SE2_0, l_cum[0], s) + + T_g0_L0 = Tangent_gn_SE2(xi_SE2_0, l_cum[0], l_cum[1]) + T_g0_s = Tangent_gn_SE2(xi_SE2_0, l_cum[0], s) + + mat_0_L0 = Ad_g0_inv_L0 @ T_g0_L0 + mat_0_s = Ad_g0_inv_s @ T_g0_s + + J_0_L0 = jnp.concatenate( + [mat_0_L0[None, :, :], jnp.zeros((num_segments - 1, 3, 3))], axis=0 + ) + J_0_s = jnp.concatenate( + [mat_0_s[None, :, :], jnp.zeros((num_segments - 1, 3, 3))], axis=0 + ) + + tuple_J_0 = (J_0_L0, J_0_s) + + # Iteration function + def J_i(tuple_J_prev: Array, i: int) -> Tuple[Tuple[Array, Array], Array]: + J_prev_Lprev, _ = tuple_J_prev + + start_index = 3 * i + xi_SE2_i = lax.dynamic_slice(xi, (start_index,), (3,)) + + Ad_gi_inv_Li = Adjoint_gn_SE2_inv(xi_SE2_i, l_cum[i], l_cum[i + 1]) + Ad_gi_inv_s = Adjoint_gn_SE2_inv(xi_SE2_i, l_cum[i], s) + + T_gi_Li = Tangent_gn_SE2(xi_SE2_i, l_cum[i], l_cum[i + 1]) + T_gi_s = Tangent_gn_SE2(xi_SE2_i, l_cum[i], s) + + mat_i_Li = Ad_gi_inv_Li @ T_gi_Li + mat_i_s = Ad_gi_inv_s @ T_gi_s + + J_new_s = lax.dynamic_update_slice( + jnp.einsum("ij, njk->nik", Ad_gi_inv_s, J_prev_Lprev), + mat_i_s[jnp.newaxis, ...], + (i, 0, 0), + ) + J_new_Li = lax.dynamic_update_slice( + jnp.einsum("ij, njk->nik", Ad_gi_inv_Li, J_prev_Lprev), + mat_i_Li[jnp.newaxis, ...], + (i, 0, 0), + ) + + tuple_J_new = (J_new_Li, J_new_s) + + return tuple_J_new, J_new_s # We accumulate J_new_s + + _, J_array = lax.scan(f=J_i, init=tuple_J_0, xs=jnp.arange(1, num_segments)) + + # Add the initial condition to the Jacobian array + J_array = jnp.concatenate([J_0_s[jnp.newaxis, ...], J_array], axis=0) + + # Extract the Jacobian for the segment that contains the point s + J_segment_SE2_local = lax.dynamic_index_in_dim( + J_array, segment_idx, axis=0, keepdims=False + ) + + # Reorder the lines to match the forward kinematics function + J_local = J_segment_SE2_local[ + :, REORDERED_LINES_FWD_KINE, : + ] # shape: (n_segments, 3, 3) + + J_local = blk_concat(J_local) # shape: (n_segments*3, 3) + J_local = J_local @ B_xi + + return J_local + + def J_explicit_global( + params: Dict[str, Array], xi: Array, s: Array, eps: float = global_eps + ) -> Array: + """ + Compute the inertial-frame jacobian of the forward kinematics function with respect to the strain vector + at a given point s using explicit expression in SE(2). + + Args: + params (Dict[str, Array]): dictionary of robot parameters. + xi (Array): strain vector of the robot. + s (Array): point coordinate along the robot in the interval [0, L_tot]. + eps (float, optional): small number to avoid singularities. Defaults to global_eps = 1e-8. + + Returns: + J_global (Array): inertial-frame jacobian of the forward kinematics function with respect to the strain vector + using explicit expression in SE(2). + """ + + # Classify the point along the robot to the corresponding segment + segment_idx, _, l_cum = classify_segment(params, s) + + xi = apply_eps_to_bend_strains(xi, eps) # Apply eps to the bending strain + + # Initial condition + xi_SE2_0 = xi[0:3] + + Ad_g0_inv_L0 = Adjoint_gn_SE2_inv(xi_SE2_0, l_cum[0], l_cum[1]) + Ad_g0_inv_s = Adjoint_gn_SE2_inv(xi_SE2_0, l_cum[0], s) + + T_g0_L0 = Tangent_gn_SE2(xi_SE2_0, l_cum[0], l_cum[1]) + T_g0_s = Tangent_gn_SE2(xi_SE2_0, l_cum[0], s) + + mat_0_L0 = Ad_g0_inv_L0 @ T_g0_L0 + mat_0_s = Ad_g0_inv_s @ T_g0_s + + J_0_L0 = jnp.concatenate( + [mat_0_L0[None, :, :], jnp.zeros((num_segments - 1, 3, 3))], axis=0 + ) + J_0_s = jnp.concatenate( + [mat_0_s[None, :, :], jnp.zeros((num_segments - 1, 3, 3))], axis=0 + ) + + tuple_J_0 = (J_0_L0, J_0_s) + + # Iteration function + def J_i(tuple_J_prev: Array, i: int) -> Tuple[Tuple[Array, Array], Array]: + J_prev_Lprev, _ = tuple_J_prev + + start_index = 3 * i + xi_SE2_i = lax.dynamic_slice(xi, (start_index,), (3,)) + + Ad_gi_inv_Li = Adjoint_gn_SE2_inv(xi_SE2_i, l_cum[i], l_cum[i + 1]) + Ad_gi_inv_s = Adjoint_gn_SE2_inv(xi_SE2_i, l_cum[i], s) + + T_gi_Li = Tangent_gn_SE2(xi_SE2_i, l_cum[i], l_cum[i + 1]) + T_gi_s = Tangent_gn_SE2(xi_SE2_i, l_cum[i], s) + + mat_i_Li = Ad_gi_inv_Li @ T_gi_Li + mat_i_s = Ad_gi_inv_s @ T_gi_s + + J_new_s = lax.dynamic_update_slice( + jnp.einsum("ij, njk->nik", Ad_gi_inv_s, J_prev_Lprev), + mat_i_s[jnp.newaxis, ...], + (i, 0, 0), + ) + J_new_Li = lax.dynamic_update_slice( + jnp.einsum("ij, njk->nik", Ad_gi_inv_Li, J_prev_Lprev), + mat_i_Li[jnp.newaxis, ...], + (i, 0, 0), + ) + + tuple_J_new = (J_new_Li, J_new_s) + + return tuple_J_new, J_new_s # We accumulate J_new_s + + _, J_array = lax.scan(f=J_i, init=tuple_J_0, xs=jnp.arange(1, num_segments)) + + # Add the initial condition to the Jacobian array + J_array = jnp.concatenate([J_0_s[jnp.newaxis, ...], J_array], axis=0) + + # Extract the Jacobian for the segment that contains the point s + J_segment_SE2_local = lax.dynamic_index_in_dim( + J_array, segment_idx, axis=0, keepdims=False + ) + + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # From local to global frame : applying the rotation of the pose at point s + + # Get the pose at point s + _, _, theta = chi_fn(params, xi, s, eps) + # Convert the pose to SE(3) representation + c, s = jnp.cos(theta), jnp.sin(theta) + R = jnp.stack([jnp.stack([c, -s]), jnp.stack([s, c])]) + g_s = jnp.block([[R, jnp.zeros((2, 1))], [jnp.zeros((1, 2)), jnp.eye(1)]]) + Adjoint_g_s = Adjoint_g_SE2(g_s) + # For each segment, compute the Jacobian in SE(3) coordinates in global frame J_i_global = Adjoint_g_s @ J_i_local + J_segment_SE2_global = jnp.einsum( + "ij,njk->nik", Adjoint_g_s, J_segment_SE2_local + ) # shape: (n_segments, 6, 6) + + # Reorder the lines to match the forward kinematics function + J_global = J_segment_SE2_global[ + :, REORDERED_LINES_FWD_KINE, : + ] # shape: (n_segments, 3, 3) + + J_global = blk_concat(J_global) # shape: (n_segments*3, 3) + J_global = J_global @ B_xi + + return J_global + + def J_Jd_autodiff( + params: Dict[str, Array], + xi: Array, + xi_d: Array, + s: Array, + eps: float = global_eps, + ) -> Tuple[Array, Array]: + """ + Compute the inertial-frame jacobian of the forward kinematics function and its time-derivative using autodiff. + + Args: + params (Dict[str, Array]): dictionary of robot parameters. + xi (Array): strain vector of the robot. + xi_d (Array): velocity vector of the robot. + s (Array): point coordinate along the robot in the interval [0, L]. + + Returns: + J (Array): inertial-frame jacobian of the forward kinematics function. + J_d (Array): inertial-frame time-derivative of the Jacobian of the forward kinematics function. + """ + + xi = apply_eps_to_bend_strains(xi, eps) # Apply eps to the bending strain + + # Compute the Jacobian of chi_fn with respect to xi + J = jacobian(lambda _xi: chi_fn(params, _xi, s, eps))(xi) + + dJ_dxi = jacobian(J)(xi) + J_d = jnp.tensordot(dJ_dxi, xi_d, axes=([2], [0])) + + # apply the strain basis to the Jacobian + J = J @ B_xi + + # apply the strain basis to the time-derivative of the Jacobian + J_d = J_d @ B_xi + + return J, J_d + + def J_Jd_explicit_local( + params: Dict[str, Array], + xi: Array, + xi_d: Array, + s: Array, + eps: float = global_eps, + ) -> Tuple[Array, Array]: + """ + Compute the body-frame jacobian and its derivative with respect to the strain vector + at a given point s using explicit expression in SE(2). + + Args: + params (Dict[str, Array]): dictionary of robot parameters. + xi (Array): strain vector of the robot. + xi_d (Array): velocity vector of the robot. + s (Array): point coordinate along the robot in the interval [0, L]. + eps (float, optional): small number to avoid singularities. Defaults to global_eps = 1e-8. + + Returns: + J_local (Array): body-frame jacobian of the forward kinematics function. + J_d_local (Array): body-frame time-derivative of the jacobian of the forward kinematics function. + """ + # Classify the point along the robot to the corresponding segment + segment_idx, _, l_cum = classify_segment(params, s) + + xi = apply_eps_to_bend_strains(xi, eps) # Apply eps to the bending strain + + # ================================= + # Computation of the Jacobian + + # Initial condition + xi_SE2_0 = xi[0:3] + + Ad_g0_inv_L0 = Adjoint_gn_SE2_inv(xi_SE2_0, l_cum[0], l_cum[1]) + Ad_g0_inv_s = Adjoint_gn_SE2_inv(xi_SE2_0, l_cum[0], s) + + T_g0_L0 = Tangent_gn_SE2(xi_SE2_0, l_cum[0], l_cum[1]) + T_g0_s = Tangent_gn_SE2(xi_SE2_0, l_cum[0], s) + + mat_0_L0 = Ad_g0_inv_L0 @ T_g0_L0 + mat_0_s = Ad_g0_inv_s @ T_g0_s + + J_0_L0 = jnp.concatenate( + [mat_0_L0[None, :, :], jnp.zeros((num_segments - 1, 3, 3))], axis=0 + ) + J_0_s = jnp.concatenate( + [mat_0_s[None, :, :], jnp.zeros((num_segments - 1, 3, 3))], axis=0 + ) + + tuple_J_0 = (J_0_L0, J_0_s) + + # Iteration function + def J_i(tuple_J_prev: Array, i: int) -> Tuple[Tuple[Array, Array], Array]: + J_prev_Lprev, _ = tuple_J_prev + + start_index = 3 * i + xi_SE2_i = lax.dynamic_slice(xi, (start_index,), (3,)) + + Ad_gi_inv_Li = Adjoint_gn_SE2_inv(xi_SE2_i, l_cum[i], l_cum[i + 1]) + Ad_gi_inv_s = Adjoint_gn_SE2_inv(xi_SE2_i, l_cum[i], s) + + T_gi_Li = Tangent_gn_SE2(xi_SE2_i, l_cum[i], l_cum[i + 1]) + T_gi_s = Tangent_gn_SE2(xi_SE2_i, l_cum[i], s) + + mat_i_Li = Ad_gi_inv_Li @ T_gi_Li + mat_i_s = Ad_gi_inv_s @ T_gi_s + + J_new_s = lax.dynamic_update_slice( + jnp.einsum("ij, njk->nik", Ad_gi_inv_s, J_prev_Lprev), + mat_i_s[jnp.newaxis, ...], + (i, 0, 0), + ) + J_new_Li = lax.dynamic_update_slice( + jnp.einsum("ij, njk->nik", Ad_gi_inv_Li, J_prev_Lprev), + mat_i_Li[jnp.newaxis, ...], + (i, 0, 0), + ) + + tuple_J_new = (J_new_Li, J_new_s) + + return tuple_J_new, J_new_s # We accumulate J_new_s + + _, J_array = lax.scan(f=J_i, init=tuple_J_0, xs=jnp.arange(1, num_segments)) + + # Add the initial condition to the Jacobian array + J_array = jnp.concatenate([J_0_s[jnp.newaxis, ...], J_array], axis=0) + + # Extract the Jacobian for the segment that contains the point s + J_segment_SE2_local = lax.dynamic_index_in_dim( + J_array, segment_idx, axis=0, keepdims=False + ) + + # ================================= + # Computation of the time-derivative of the Jacobian + + idx_range = jnp.arange(num_segments) + xi_d_SE2_i = vmap(lambda i: lax.dynamic_slice(xi_d, (3 * i,), (3,)))( + idx_range + ) # shape: (num_segments, 3) + S_i = vmap( + lambda i: lax.dynamic_index_in_dim( + J_segment_SE2_local, i, axis=0, keepdims=False + ) + )(idx_range) # shape: (num_segments, 3, 3) + sum_Sj_xi_d_j = compute_weighted_sums( + J_segment_SE2_local, xi_d_SE2_i, num_segments + ) # shape: (num_segments, 3) + adjoint_sum = vmap(adjoint_SE2)(sum_Sj_xi_d_j) # shape: (num_segments, 3, 3) + + # Compute the time-derivative of the Jacobian + J_d_segment_SE2_local = jnp.einsum( + "ijk, ikl->ijl", adjoint_sum, S_i + ) # shape: (num_segments, 3, 3) + + # Replace the elements of J_d_segment_SE2 for i > segment_idx by null matrices + J_d_segment_SE2_local = jnp.where( + jnp.arange(num_segments)[:, None, None] > segment_idx, + jnp.zeros_like(J_d_segment_SE2_local), + J_d_segment_SE2_local, + ) + + # Reorder the lines to match the forward kinematics function + J_local = J_segment_SE2_local[ + :, REORDERED_LINES_FWD_KINE, : + ] # shape: (n_segments, 3, 3) + J_d_local = J_d_segment_SE2_local[ + :, REORDERED_LINES_FWD_KINE, : + ] # shape: (n_segments, 3, 3) + + # Concatenate the Jacobians and their time-derivatives for all segments + J_local = blk_concat(J_local) # shape: (n_segments*3, 3) + J_d_local = blk_concat(J_d_local) # shape: (n_segments*3, 3) + + # Apply the strain basis to the Jacobian and its time-derivative + J_local = J_local @ B_xi + J_d_local = J_d_local @ B_xi + + return J_local, J_d_local + + def J_Jd_explicit_global( + params: Dict[str, Array], + xi: Array, + xi_d: Array, + s: Array, + eps: float = global_eps, + ) -> Tuple[Array, Array]: + """ + Compute the inertial-frame jacobian and its derivative with respect to the strain vector + at a given point s using explicit expression in SE(2). + + Args: + params (Dict[str, Array]): dictionary of robot parameters. + xi (Array): strain vector of the robot. + xi_d (Array): velocity vector of the robot. + s (Array): point coordinate along the robot in the interval [0, L]. + eps (float, optional): small number to avoid singularities. Defaults to global_eps = 1e-8. + + Returns: + J_global (Array): inertial-frame jacobian of the forward kinematics function. + J_d_global (Array): inertial-frame time-derivative of the jacobian of the forward kinematics function. + """ + # Classify the point along the robot to the corresponding segment + segment_idx, _, l_cum = classify_segment(params, s) + + xi = apply_eps_to_bend_strains(xi, eps) # Apply eps to the bending strain + + # ================================= + # Computation of the Jacobian + + # Initial condition + xi_SE2_0 = xi[0:3] + + Ad_g0_inv_L0 = Adjoint_gn_SE2_inv(xi_SE2_0, l_cum[0], l_cum[1]) + Ad_g0_inv_s = Adjoint_gn_SE2_inv(xi_SE2_0, l_cum[0], s) + + T_g0_L0 = Tangent_gn_SE2(xi_SE2_0, l_cum[0], l_cum[1]) + T_g0_s = Tangent_gn_SE2(xi_SE2_0, l_cum[0], s) + + mat_0_L0 = Ad_g0_inv_L0 @ T_g0_L0 + mat_0_s = Ad_g0_inv_s @ T_g0_s + + J_0_L0 = jnp.concatenate( + [mat_0_L0[None, :, :], jnp.zeros((num_segments - 1, 3, 3))], axis=0 + ) + J_0_s = jnp.concatenate( + [mat_0_s[None, :, :], jnp.zeros((num_segments - 1, 3, 3))], axis=0 + ) + + tuple_J_0 = (J_0_L0, J_0_s) + + # Iteration function + def J_i(tuple_J_prev: Array, i: int) -> Tuple[Tuple[Array, Array], Array]: + J_prev_Lprev, _ = tuple_J_prev + + start_index = 3 * i + xi_SE2_i = lax.dynamic_slice(xi, (start_index,), (3,)) + + Ad_gi_inv_Li = Adjoint_gn_SE2_inv(xi_SE2_i, l_cum[i], l_cum[i + 1]) + Ad_gi_inv_s = Adjoint_gn_SE2_inv(xi_SE2_i, l_cum[i], s) + + T_gi_Li = Tangent_gn_SE2(xi_SE2_i, l_cum[i], l_cum[i + 1]) + T_gi_s = Tangent_gn_SE2(xi_SE2_i, l_cum[i], s) + + mat_i_Li = Ad_gi_inv_Li @ T_gi_Li + mat_i_s = Ad_gi_inv_s @ T_gi_s + + J_new_s = lax.dynamic_update_slice( + jnp.einsum("ij, njk->nik", Ad_gi_inv_s, J_prev_Lprev), + mat_i_s[jnp.newaxis, ...], + (i, 0, 0), + ) + J_new_Li = lax.dynamic_update_slice( + jnp.einsum("ij, njk->nik", Ad_gi_inv_Li, J_prev_Lprev), + mat_i_Li[jnp.newaxis, ...], + (i, 0, 0), + ) + + tuple_J_new = (J_new_Li, J_new_s) + + return tuple_J_new, J_new_s # We accumulate J_new_s + + _, J_array = lax.scan(f=J_i, init=tuple_J_0, xs=jnp.arange(1, num_segments)) + + # Add the initial condition to the Jacobian array + J_array = jnp.concatenate([J_0_s[jnp.newaxis, ...], J_array], axis=0) + + # Extract the Jacobian for the segment that contains the point s + J_segment_SE2_local = lax.dynamic_index_in_dim( + J_array, segment_idx, axis=0, keepdims=False + ) + + # ================================= + # Computation of the time-derivative of the Jacobian + + idx_range = jnp.arange(num_segments) + xi_d_SE2_i = vmap(lambda i: lax.dynamic_slice(xi_d, (3 * i,), (3,)))( + idx_range + ) # shape: (num_segments, 3) + S_i = vmap( + lambda i: lax.dynamic_index_in_dim( + J_segment_SE2_local, i, axis=0, keepdims=False + ) + )(idx_range) # shape: (num_segments, 3, 3) + sum_Sj_xi_d_j = compute_weighted_sums( + J_segment_SE2_local, xi_d_SE2_i, num_segments + ) # shape: (num_segments, 3) + adjoint_sum = vmap(adjoint_SE2)(sum_Sj_xi_d_j) # shape: (num_segments, 3, 3) + + # Compute the time-derivative of the Jacobian + J_d_segment_SE2_local = jnp.einsum( + "ijk, ikl->ijl", adjoint_sum, S_i + ) # shape: (num_segments, 3, 3) + + # Replace the elements of J_d_segment_SE2 for i > segment_idx by null matrices + J_d_segment_SE2_local = jnp.where( + jnp.arange(num_segments)[:, None, None] > segment_idx, + jnp.zeros_like(J_d_segment_SE2_local), + J_d_segment_SE2_local, + ) + + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # From local to global frame : applying the rotation of the pose at point s + + # Get the pose at point s + _, _, theta = chi_fn(params, xi, s, eps) + # Convert the pose to SE(3) representation + c, s = jnp.cos(theta), jnp.sin(theta) + R = jnp.stack([jnp.stack([c, -s]), jnp.stack([s, c])]) + g_s = jnp.block([[R, jnp.zeros((2, 1))], [jnp.zeros((1, 2)), jnp.eye(1)]]) + Adjoint_g_s = Adjoint_g_SE2(g_s) + # For each segment, compute the Jacobian in SE(3) coordinates in global frame J_i_global = Adjoint_g_s @ J_i_local + J_segment_SE2_global = jnp.einsum( + "ij,njk->nik", Adjoint_g_s, J_segment_SE2_local + ) # shape: (n_segments, 6, 6) + J_d_segment_global = jnp.einsum( + "ij,njk->nik", Adjoint_g_s, J_d_segment_SE2_local + ) # shape: (n_segments, 6, 6) + + # Reorder the lines to match the forward kinematics function + J_global = J_segment_SE2_global[ + :, REORDERED_LINES_FWD_KINE, : + ] # shape: (n_segments, 3, 3) + J_d_global = J_d_segment_global[ + :, REORDERED_LINES_FWD_KINE, : + ] # shape: (n_segments, 3, 3) + + # Concatenate the Jacobians and their time-derivatives for all segments + J_global = blk_concat(J_global) # shape: (n_segments*3, 3) + J_d_global = blk_concat(J_d_global) # shape: (n_segments*3, 3) + + # Apply the strain basis to the Jacobian and its time-derivative + J_global = J_global @ B_xi + J_d_global = J_d_global @ B_xi + + return J_global, J_d_global + + if jacobian_type == "explicit": + jacobian_fn_xi = J_explicit_global + J_Jd = J_Jd_explicit_global + elif jacobian_type == "autodiff": + jacobian_fn_xi = J_autodiff + J_Jd = J_Jd_autodiff + + def jacobian_fn( + params: Dict[str, Array], q: Array, s: Array, eps: float = global_eps + ) -> Array: + """ + Compute the inertial-frame jacobian of the forward kinematics function with respect to the configuration vector q. + + Args: + params (Dict[str, Array]): dictionary of robot parameters. + q (Array): configuration vector of the robot. + s (Array): point coordinate along the robot in the interval [0, L]. + eps (float, optional): small number to avoid singularities. Defaults to global_eps = 1e-8. + + Returns: + J (Array): inertial-frame jacobian of the forward kinematics function with respect to the strain vector. + """ + # Map the configuration to the strains + xi = xi_eq + B_xi @ q + + # Add a small number to the bending strain to avoid singularities + xi = apply_eps_to_bend_strains(xi, eps) + + # Compute the Jacobian of chi_fn with respect to xi + J = jacobian_fn_xi(params, xi, s) + return J + + def B_autodiff_fn(params: Dict[str, Array], xi: Array) -> Array: + """ + Compute the mass / inertia matrix of the robot using autodiff. + + Args: + params (Dict[str, Array]): dictionary of robot parameters. + xi (Array): strain vector of the robot. + + Returns: + B (Array): mass / inertia matrix of the robot. + """ + # Extract the parameters + rho = params["rho"] # density of each segment [kg/m^3] + l = params["l"] # length of each segment [m] + r = params["r"] # radius of each segment [m] + + # Usefull derived quantities + A = jnp.pi * r**2 # cross-sectional area of each segment [m^2] + Ib = A**2 / (4 * jnp.pi) # second moment of area of each segment [m^4] + + l_cum = jnp.cumsum(jnp.concatenate([jnp.zeros(1), l])) + + # Compute each integral + def compute_integral(i): + if integration_type == "gauss-legendre": + Xs, Ws, nGauss = gauss_quadrature( + N_GQ=param_integration, a=l_cum[i], b=l_cum[i + 1] + ) + + J_all = vmap(lambda s: J_autodiff(params, xi, s))(Xs) + Jp_all = J_all[:, :2, :] + Jo_all = J_all[:, 2:, :] + + integrand_JpT_Jp = jnp.einsum("nij,nik->njk", Jp_all, Jp_all) + integrand_JoT_Jo = jnp.einsum("nij,nik->njk", Jo_all, Jo_all) + + integral_Jp = jnp.sum(Ws[:, None, None] * integrand_JpT_Jp, axis=0) + integral_Jo = jnp.sum(Ws[:, None, None] * integrand_JoT_Jo, axis=0) + + integral = rho[i] * A[i] * integral_Jp + rho[i] * Ib[i] * integral_Jo + + elif integration_type == "gauss-kronrad": + rule = GaussKronrodRule(order=param_integration) + + def integrand(s): + J = J_autodiff(params, xi, s) + Jp = J[:2, :] + Jo = J[2:, :] + return rho[i] * A[i] * Jp.T @ Jp + rho[i] * Ib[i] * Jo.T @ Jo + + integral, _, _, _ = rule.integrate( + integrand, l_cum[i], l_cum[i + 1], args=() + ) + + elif integration_type == "trapezoid": + xs = jnp.linspace(l_cum[i], l_cum[i + 1], param_integration) + + J_all = vmap(lambda s: J_autodiff(params, xi, s))(xs) + Jp_all = J_all[:, :2, :] + Jo_all = J_all[:, 2:, :] + + integrand_JpT_Jp = jnp.einsum("nij,nik->njk", Jp_all, Jp_all) + integrand_JoT_Jo = jnp.einsum("nij,nik->njk", Jo_all, Jo_all) + + integral_Jp = jscipy.integrate.trapezoid(integrand_JpT_Jp, x=xs, axis=0) + integral_Jo = jscipy.integrate.trapezoid(integrand_JoT_Jo, x=xs, axis=0) + + integral = rho[i] * A[i] * integral_Jp + rho[i] * Ib[i] * integral_Jo + + return integral + + # Compute the cumulative integral + indices = jnp.arange(num_segments) + integrals = vmap(compute_integral)(indices) + + B = jnp.sum(integrals, axis=0) + + return B + + def C_autodiff_fn( + params: Dict[str, Array], xi: Array, xi_d: Array + ) -> Tuple[Array, Array]: + """ + Compute the Coriolis / centrifugal matrix of the robot + using autodiff. + + Args: + params (Dict[str, Array]): dictionary of robot parameters. + xi (Array): strain vector of the robot. + xi_d (Array): velocity vector of the robot. + + Returns: + C (Array): Coriolis / centrifugal matrix of the robot. + """ + + n_xi = 3 * num_segments + + def christoffel_fn(i, j, k): + dB_ij = grad(lambda x: B_autodiff_fn(params, x)[i, j])(xi)[k] + dB_ik = grad(lambda x: B_autodiff_fn(params, x)[i, k])(xi)[j] + dB_jk = grad(lambda x: B_autodiff_fn(params, x)[j, k])(xi)[i] + return 0.5 * (dB_ij + dB_ik - dB_jk) + + # # =========================================== + # # Native for loop version + # C = jnp.zeros((n_xi, n_xi)) # Initialize the Coriolis matrix + # for i in range(n_xi): + # for j in range(n_xi): + # for k in range(n_xi): + # christoffel_symbol = christoffel_fn(i, j, k) + # C = C.at[i,j].add(christoffel_symbol * xi_d[k]) + + # # =========================================== + # # For loop with LAX version + # C = jnp.zeros((n_xi, n_xi)) # Initialize the Coriolis matrix + # def body_i(i, C): + # def body_j(j, C): + # def body_k(k, acc): + # christoffel_symbol = christoffel_fn(i, j, k) + # coeff = christoffel_symbol * xi_d[k] + # return acc + coeff + # C_ij = lax.fori_loop(0, n_xi, body_k, 0.0) + # return C.at[i, j].set(C_ij) + # return lax.fori_loop(0, n_xi, body_j, C) + # C = lax.fori_loop(0, n_xi, body_i, C) + + # =========================================== + # Vmap version + def C_ij(i, j): + cs_k = vmap(lambda k: christoffel_fn(i, j, k))(jnp.arange(n_xi)) + return jnp.dot(cs_k, xi_d) + + C = vmap(lambda i: vmap(lambda j: C_ij(i, j))(jnp.arange(n_xi)))( + jnp.arange(n_xi) + ) + + # # =========================================== + # # LAX map version + # def C_ij(i, j, xi_d, n_xi): + # cs_k = lax.map(lambda k: christoffel_fn(i, j, k), jnp.arange(n_xi)) + # return jnp.dot(cs_k, xi_d) + # C = jnp.stack( + # jnp.stack( + # lax.map( + # lambda i: lax.map( + # lambda j: C_ij(i, j, xi_d, n_xi), + # jnp.arange(n_xi) + # ), + # jnp.arange(n_xi)) + # ) + # ) + + return C + + def B_explicit_fn(params: Dict[str, Array], xi: Array) -> Array: + """ + Compute the mass / inertia matrix of the robot using explicit expression. + + Args: + params (Dict[str, Array]): dictionary of robot parameters. + xi (Array): strain vector of the robot. + + Returns: + B (Array): mass / inertia matrix of the robot. + """ + # Extract the parameters + rho = params["rho"] # density of each segment [kg/m^3] + l = params["l"] # length of each segment [m] + r = params["r"] # radius of each segment [m] + + # Usefull derived quantities + A = jnp.pi * r**2 # cross-sectional area of each segment [m^2] + Ib = A**2 / (4 * jnp.pi) # second moment of area of each segment [m^4] + + l_cum = jnp.cumsum(jnp.concatenate([jnp.zeros(1), l])) + + # Compute each integral + def compute_integral(i): + if integration_type == "gauss-legendre": + Xs, Ws, nGauss = gauss_quadrature( + N_GQ=param_integration, a=l_cum[i], b=l_cum[i + 1] + ) + + J_all = vmap(lambda s: J_explicit_local(params, xi, s))(Xs) + Jp_all = J_all[:, :2, :] + Jo_all = J_all[:, 2:, :] + + integrand_JpT_Jp = jnp.einsum("nij,nik->njk", Jp_all, Jp_all) + integrand_JoT_Jo = jnp.einsum("nij,nik->njk", Jo_all, Jo_all) + + integral_Jp = jnp.sum(Ws[:, None, None] * integrand_JpT_Jp, axis=0) + integral_Jo = jnp.sum(Ws[:, None, None] * integrand_JoT_Jo, axis=0) + + integral = rho[i] * A[i] * integral_Jp + rho[i] * Ib[i] * integral_Jo + + elif integration_type == "gauss-kronrad": + rule = GaussKronrodRule(order=param_integration) + + def integrand(s): + J = J_explicit_local(params, xi, s) + Jp = J[:2, :] + Jo = J[2:, :] + return rho[i] * A[i] * Jp.T @ Jp + rho[i] * Ib[i] * Jo.T @ Jo + + integral, _, _, _ = rule.integrate( + integrand, l_cum[i], l_cum[i + 1], args=() + ) + + elif integration_type == "trapezoid": + xs = jnp.linspace(l_cum[i], l_cum[i + 1], param_integration) + + J_all = vmap(lambda s: J_explicit_local(params, xi, s))(xs) + Jp_all = J_all[:, :2, :] + Jo_all = J_all[:, 2:, :] + + integrand_JpT_Jp = jnp.einsum("nij,nik->njk", Jp_all, Jp_all) + integrand_JoT_Jo = jnp.einsum("nij,nik->njk", Jo_all, Jo_all) + + integral_Jp = jscipy.integrate.trapezoid(integrand_JpT_Jp, x=xs, axis=0) + integral_Jo = jscipy.integrate.trapezoid(integrand_JoT_Jo, x=xs, axis=0) + + integral = rho[i] * A[i] * integral_Jp + rho[i] * Ib[i] * integral_Jo + + return integral + + # Compute the cumulative integral + indices = jnp.arange(num_segments) + integrals = vmap(compute_integral)(indices) + + B = jnp.sum(integrals, axis=0) + + return B + + def C_explicit_fn( + params: Dict[str, Array], xi: Array, xi_d: Array + ) -> Tuple[Array, Array]: + """ + Compute the Coriolis / centrifugal matrix of the robot using explicit expression. + + Args: + params (Dict[str, Array]): dictionary of robot parameters. + xi (Array): strain vector of the robot. + xi_d (Array): velocity vector of the robot. + + Returns: + C (Array): Coriolis / centrifugal matrix of the robot. + """ + + # Extract the parameters + rho = params["rho"] # density of each segment [kg/m^3] + l = params["l"] # length of each segment [m] + r = params["r"] # radius of each segment [m] + + # Usefull derived quantities + A = jnp.pi * r**2 # cross-sectional area of each segment [m^2] + Ib = A**2 / (4 * jnp.pi) # second moment of area of each segment [m^4] + + l_cum = jnp.cumsum(jnp.concatenate([jnp.zeros(1), l])) + + # Compute each integral + def compute_integral_C(i): + if integration_type == "gauss-legendre": + Xs, Ws, nGauss = gauss_quadrature( + N_GQ=param_integration, a=l_cum[i], b=l_cum[i + 1] + ) + + J_all, J_d_all = vmap( + lambda s: J_Jd_explicit_local(params, xi, xi_d, s) + )(Xs) # [[Jp1],[Jp2],[Jo]] + J_all = J_all[:, REORDERED_LINES_STRAIN, :] # [[Jo],[Jp1],[Jp2]] + J_d_all = J_d_all[ + :, REORDERED_LINES_STRAIN, : + ] # [[Jo_d],[Jp1_d],[Jp2_d]] + M_a = rho[i] * jnp.diag( + jnp.stack([Ib[i], A[i], A[i]], axis=0) + ) # [[O],[P],[P]] + + integrand_C = vmap( + lambda _J_i, _J_d_i: ( + _J_i.T + @ (adjoint_star_SE2((_J_i @ xi_d)) @ M_a @ _J_i + M_a @ _J_d_i) + ) + )(J_all, J_d_all) + + integral_C = jnp.sum(Ws[:, None, None] * integrand_C, axis=0) + + elif integration_type == "gauss-kronrad": + rule = GaussKronrodRule(order=param_integration) + + def integrand(s): + J, J_d = J_Jd_explicit_local( + params, xi, xi_d, s + ) # [[Jp1],[Jp2],[Jo]] + J = J[REORDERED_LINES_STRAIN, :] # [Jo, Jp1, Jp2] + J_d = J_d[REORDERED_LINES_STRAIN, :] # [Jo_d, Jp1_d, Jp2_d] + M_a = rho[i] * jnp.diag( + jnp.stack([Ib[i], A[i], A[i]], axis=0) + ) # [[O],[P],[P]] + return J.T @ (adjoint_star_SE2((J @ xi_d)) @ M_a @ J + M_a @ J_d) + + integral_C, _, _, _ = rule.integrate( + integrand, l_cum[i], l_cum[i + 1], args=() + ) + + elif integration_type == "trapezoid": + xs = jnp.linspace(l_cum[i], l_cum[i + 1], param_integration) + + J_all, J_d_all = vmap( + lambda s: J_Jd_explicit_local(params, xi, xi_d, s) + )(xs) # [[Jp1],[Jp2],[Jo]] + J_all = J_all[:, REORDERED_LINES_STRAIN, :] # [[Jo],[Jp1],[Jp2]] + J_d_all = J_d_all[ + :, REORDERED_LINES_STRAIN, : + ] # [[Jo_d],[Jp1_d],[Jp2_d]] + M_a = rho[i] * jnp.diag( + jnp.stack([Ib[i], A[i], A[i]], axis=0) + ) # [[O],[P],[P]] + + integrand_C = vmap( + lambda _J_i, _J_d_i: ( + _J_i.T + @ (adjoint_star_SE2((_J_i @ xi_d)) @ M_a @ _J_i + M_a @ _J_d_i) + ) + )(J_all, J_d_all) + + integral_C = jscipy.integrate.trapezoid(integrand_C, x=xs, axis=0) + + return integral_C + + # Compute the cumulative integral + indices = jnp.arange(num_segments) + integrals = vmap(compute_integral_C)(indices) + + C = jnp.sum(integrals, axis=0) + + return C + + def U_g_fn_xi( + params: Dict[str, Array], xi: Array, eps: float = global_eps + ) -> Array: + """ + Compute the gravity vector of the robot. + + Args: + params (Dict[str, Array]): dictionary of robot parameters. + xi (Array): strain vector of the robot. + + Returns: + U_g (Array): gravity vector of the robot. + """ + + # Extract the parameters + g = params["g"] # gravity vector [m/s^2] + rho = params["rho"] # density of each segment [kg/m^3] + l = params["l"] # length of each segment [m] + r = params["r"] # radius of each segment [m] + + # Usefull derived quantitie + A = jnp.pi * r**2 # cross-sectional area of each segment [m^2] + + l_cum = jnp.cumsum(jnp.concatenate([jnp.zeros(1), l])) + + # Compute each integral + def compute_integral(i): + if integration_type == "gauss-legendre": + Xs, Ws, nGauss = gauss_quadrature( + N_GQ=param_integration, a=l_cum[i], b=l_cum[i + 1] + ) + chi_s = vmap(lambda s: chi_fn(params, xi, s, eps))(Xs) + p_s = chi_s[:, :2] + integrand = -rho[i] * A[i] * jnp.einsum("ij,j->i", p_s, g) + + # Compute the integral + integral = jnp.sum(Ws * integrand) + + elif integration_type == "gauss-kronrad": + rule = GaussKronrodRule(order=param_integration) + + def integrand(s): + chi_s = chi_fn(params, xi, s, eps) + p_s = chi_s[:2] + return -rho[i] * A[i] * jnp.dot(p_s, g) + + # Compute the integral + integral, _, _, _ = rule.integrate( + integrand, l_cum[i], l_cum[i + 1], args=() + ) + + elif integration_type == "trapezoid": + xs = jnp.linspace(l_cum[i], l_cum[i + 1], param_integration) + chi_s = vmap(lambda s: chi_fn(params, xi, s, eps))(xs) + p_s = chi_s[:, :2] + integrand = -rho[i] * A[i] * jnp.einsum("ij,j->i", p_s, g) + + # Compute the integral + integral = jscipy.integrate.trapezoid(integrand, x=xs) + + return integral + + # Compute the cumulative integral + indices = jnp.arange(num_segments) + integrals = vmap(compute_integral)(indices) + + U_g = jnp.sum(integrals) + + return U_g + + def G_autodiff_fn( + params: Dict[str, Array], + xi: Array, + ) -> Array: + """ + Compute the gravity vector of the robot + using autodiff. + + Args: + params (Dict[str, Array]): dictionary of robot parameters. + xi (Array): strain vector of the robot. + + Returns: + G (Array) : gravity vector of the robot. + """ + + G = jacobian(lambda _xi: U_g_fn_xi(params, _xi))(xi) + + return G + + def G_explicit_fn( + params: Dict[str, Array], + xi: Array, + ) -> Array: + """ + Compute the gravity vector of the robot + using explicit expressions. + + Args: + params (Dict[str, Array]): dictionary of robot parameters. + xi (Array): strain vector of the robot. + + Returns: + G (Array): gravity vector of the robot. + """ + # Extract the parameters + g = params["g"] # gravity vector [m/s^2] + rho = params["rho"] # density of each segment [kg/m^3] + l = params["l"] # length of each segment [m] + r = params["r"] # radius of each segment [m] + + # Usefull derived quantitie + A = jnp.pi * r**2 # cross-sectional area of each segment [m^2] + + l_cum = jnp.cumsum(jnp.concatenate([jnp.zeros(1), l])) + + # Compute each integral + def compute_integral(i): + if integration_type == "gauss-legendre": + Xs, Ws, nGauss = gauss_quadrature( + N_GQ=param_integration, a=l_cum[i], b=l_cum[i + 1] + ) + + J_all = vmap(lambda s: J_explicit_global(params, xi, s))(Xs) + Jp_all = J_all[:, :2, :] # shape: (nGauss, n_segments, 3, 3) + + # Compute the integrand + integrand = -rho[i] * A[i] * jnp.einsum("ijk,j->ik", Jp_all, g) + + # Multiply each element of integrand by the corresponding weight + weighted_integrand = jnp.einsum("i, ij->ij", Ws, integrand) + + # Compute the integral + integral = jnp.sum( + weighted_integrand, axis=0 + ) # sum over the Gauss points + + elif integration_type == "gauss-kronrad": + rule = GaussKronrodRule(order=param_integration) + + def integrand(s): + J = J_explicit_global(params, xi, s) + Jp = J[:2, :] + return -rho[i] * A[i] * jnp.dot(g, Jp) + + # Compute the integral + integral, _, _, _ = rule.integrate( + integrand, l_cum[i], l_cum[i + 1], args=() + ) + + elif integration_type == "trapezoid": + xs = jnp.linspace(l_cum[i], l_cum[i + 1], param_integration) + + J_all = vmap(lambda s: J_explicit_global(params, xi, s))(xs) + Jp_all = J_all[:, :2, :] # shape: (nGauss, n_segments, 3, 3) + + # Compute the integrand + integrand = -rho[i] * A[i] * jnp.einsum("ijk,j->ik", Jp_all, g) + + # Compute the integral + integral = jnp.sum(integrand, axis=0) + + return integral + + # Compute the cumulative integral + indices = jnp.arange(num_segments) + integrals = vmap(compute_integral)(indices) + + G = jnp.sum(integrals, axis=0) # sum over the segments + + return G + + if jacobian_type == "explicit": + B_fn_xi = B_explicit_fn + C_fn_xi = C_explicit_fn + G_fn_xi = G_explicit_fn + elif jacobian_type == "autodiff": + B_fn_xi = B_autodiff_fn + C_fn_xi = C_autodiff_fn + G_fn_xi = G_autodiff_fn + + def dynamical_matrices_fn( + params: Dict[str, Array], q: Array, q_d: Array, eps: float = global_eps + ) -> Tuple[Array, Array, Array, Array, Array, Array]: + """ + Compute the dynamical matrices of the robot. + + Args: + params (Dict[str, Array]): dictionary of robot parameters. + q (Array): configuration vector of the robot. + q_d (Array): velocity vector of the robot. + eps (float, optional): small number to avoid singularities. Defaults to global_eps. + + Returns: + B (Array): mass / inertia matrix of the robot. (shape: (n_q, n_q)) + C (Array): Coriolis / centrifugal matrix of the robot. (shape: (n_q, n_q)) + G (Array): gravity vector of the robot. (shape: (n_q,)) + K (Array): elastic vector of the robot. (shape: (n_q,)) + D (Array): dissipative matrix of the robot. (shape: (n_q, n_q)) + alpha (Array): actuation matrix of the robot. (shape: (n_q, n_tau)) + """ + # Map the configuration to the strains + xi = xi_eq + B_xi @ q + xi_d = B_xi @ q_d + + # Add a small number to the bending strain to avoid singularities + xi_epsed = apply_eps_to_bend_strains(xi, eps) + + # Compute the stiffness matrix + K = stiffness_fn(params, B_xi, formulate_in_strain_space=True) + # Apply the strain basis to the stiffness matrix + K = B_xi.T @ K @ (xi - xi_eq) # evaluate K(xi) = K @ xi + + # Compute the actuation matrix + A = actuation_mapping_fn( + forward_kinematics_fn, jacobian_fn, params, B_xi, q, eps + ) + # Apply the strain basis to the actuation matrix + alpha = A + + # Dissipative matrix + D = params.get("D", jnp.zeros((n_xi, n_xi))) + # Apply the strain basis to the dissipative matrix + D = B_xi.T @ D @ B_xi + + # Mass/inertia matrix + B = B_xi.T @ B_fn_xi(params, xi_epsed) @ B_xi + + # Coriolis matrix + C = B_xi.T @ C_fn_xi(params, xi_epsed, xi_d) @ B_xi + + # Gravitational matrix + G = B_xi.T @ G_fn_xi(params, xi_epsed).squeeze() + + return B, C, G, K, D, alpha + + def kinetic_energy_fn( + params: Dict[str, Array], q: Array, q_d: Array, eps: float = global_eps + ) -> Array: + """ + Compute the kinetic energy of the system. + Args: + params (Dict[str, Array]): dictionary of robot parameters. + q (Array): generalized coordinates of shape (n_q, ) + q_d (Array): generalized velocities of shape (n_q, ) + eps (float, optional): small number to avoid singularities (e.g., division by zero). Defaults to global_eps. + Returns: + T (Array): kinetic energy of shape () + """ + + # Map the configuration to the strains + xi = xi_eq + B_xi @ q + # Add a small number to the bending strain to avoid singularities + xi_epsed = apply_eps_to_bend_strains(xi, eps) + + # Compute the inertia matrix + B = B_fn_xi(params, xi_epsed) + + # Kinetic energy + T = (0.5 * q_d.T @ B @ q_d).squeeze() + + return T + + def potential_energy_fn( + params: Dict[str, Array], q: Array, eps: float = global_eps + ) -> Array: + """ + Compute the potential energy of the system. + Args: + params (Dict[str, Array]): dictionary of robot parameters. + q (Array): generalized coordinates of shape (n_q, ) + eps (float, optional): small number to avoid singularities (e.g., division by zero). Defaults to global_eps. + Returns: + U (Array): potential energy of shape () + """ + # Map the configuration to the strains + xi = xi_eq + B_xi @ q + # Add a small number to the bending strain to avoid singularities + xi_epsed = apply_eps_to_bend_strains(xi, eps) + + # compute the stiffness matrix + K = stiffness_fn(params, B_xi, formulate_in_strain_space=True) + # elastic energy + U_K = 0.5 * (xi - xi_eq).T @ K @ (xi - xi_eq) # evaluate K(xi) = K @ xi + + # gravitational potential energy + U_G = U_g_fn_xi(params, xi_epsed) + + # total potential energy + U = (U_G + U_K).squeeze() + + return U + + def energy_fn( + params: Dict[str, Array], q: Array, q_d: Array, eps: float = global_eps + ) -> Array: + """ + Compute the total energy of the system. + Args: + params (Dict[str, Array]): dictionary of robot parameters. + q (Array): generalized coordinates of shape (n_q, ) + q_d (Array): generalized velocities of shape (n_q, ) + eps (float, optional): small number to avoid singularities (e.g., division by zero). Defaults to global_eps. + Returns: + E (Array): total energy of shape () + """ + T = kinetic_energy_fn(params, q, q_d, eps) + U = potential_energy_fn(params, q, eps) + E = T + U + + return E + + def operational_space_dynamical_matrices_fn( + params: Dict[str, Array], + q: Array, + q_d: Array, + s: Array, + B: Array, + C: Array, + operational_space_selector: Tuple = (True, True, True), + eps: float = global_eps, + ) -> Tuple[Array, Array, Array, Array, Array]: + """ + Compute the dynamics in operational space. + The implementation is based on Chapter 7.8 of "Modelling, Planning and Control of Robotics" by + Siciliano, Sciavicco, Villani, Oriolo. + Args: + params (Dict[str, Array]): dictionary of robot parameters. + q (Array): generalized coordinates of shape (n_q, ) + q_d (Array): generalized velocities of shape (n_q, ) + s (Array): point coordinate along the robot in the interval [0, L]. + B (Array): mass / inertia matrix of the robot in the generalized coordinates of shape (n_q, n_q) + C (Array): Coriolis / centrifugal matrix of the robot in the generalized coordinates of shape (n_q, n_q) + operational_space_selector (Tuple, optional): tuple of shape (3,) to select the operational space variables. + For example, (True, True, False) selects only the positional components of the operational space. + eps (float, optional): small number to avoid singularities (e.g., division by zero). Defaults to global_eps. + Returns: + Lambda (Array): inertia matrix in the operational space of shape (3, 3) + mu (Array): coriolis and centrifugal matrix in the operational space of shape (3, ) + J (Array): inertial-frame jacobian of the end-effector pose with respect to the generalized coordinates + J_d (Array): inertial-frame time-derivative of the jacobian of the end-effector pose with respect to the generalized coordinates + of shape (3, n_q) + JB_pinv (Array): Dynamically-consistent pseudo-inverse of the inertial-frame jacobian. Allows the mapping of torques + from the generalized coordinates to the operational space: f = JB_pinv.T @ tau_q + Shape (n_q, 3) + """ + # Map the configuration to the strains + xi = xi_eq + B_xi @ q + xi_d = B_xi @ q_d + + # classify the point along the robot to the corresponding segment + _, s_segment, _ = classify_segment(params, s) + + # make operational_space_selector a boolean array + operational_space_selector = onp.array(operational_space_selector, dtype=bool) + + # Jacobian and its time-derivative + J, J_d = J_Jd(params, xi, xi_d, s_segment, eps) + J = jnp.squeeze(J) + J_d = jnp.squeeze(J_d) + + J = J[operational_space_selector, :] + J_d = J_d[operational_space_selector, :] + + # inverse of the inertia matrix in the configuration space + B_inv = jnp.linalg.inv(B) + + Lambda = jnp.linalg.inv( + J @ B_inv @ J.T + ) # inertia matrix in the operational space + mu = Lambda @ ( + J @ B_inv @ C - J_d + ) # coriolis and centrifugal matrix in the operational space + + JB_pinv = ( + B_inv @ J.T @ Lambda + ) # dynamically-consistent pseudo-inverse of the Jacobian + + return Lambda, mu, J, J_d, JB_pinv + + auxiliary_fns: Dict[str, Callable[..., Any]] = { + "apply_eps_to_bend_strains": apply_eps_to_bend_strains, + "classify_segment": classify_segment, + "stiffness_fn": stiffness_fn, + "actuation_mapping_fn": actuation_mapping_fn, + "jacobian_fn": jacobian_fn, + "kinetic_energy_fn": kinetic_energy_fn, + "potential_energy_fn": potential_energy_fn, + "energy_fn": energy_fn, + "operational_space_dynamical_matrices_fn": operational_space_dynamical_matrices_fn, + } + + return B_xi, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns diff --git a/src/jsrm/systems/planar_pcs.py b/src/jsrm/systems/planar_pcs_sym.py similarity index 93% rename from src/jsrm/systems/planar_pcs.py rename to src/jsrm/systems/planar_pcs_sym.py index fbfebcb..64ac869 100644 --- a/src/jsrm/systems/planar_pcs.py +++ b/src/jsrm/systems/planar_pcs_sym.py @@ -54,7 +54,6 @@ def factory( # symbols for robot parameters params_syms = sym_exps["params_syms"] - @jit def select_params_for_lambdify_fn(params: Dict[str, Array]) -> List[Array]: """ Select the parameters for lambdify @@ -143,41 +142,41 @@ def select_params_for_lambdify_fn(params: Dict[str, Array]) -> List[Array]: params_syms_cat + sym_exps["state_syms"]["xi"], sym_exps["exps"]["U_g"], "jax" ) - compute_stiffness_matrix_for_all_segments_fn = vmap( - compute_planar_stiffness_matrix - ) + compute_stiffness_matrix_for_all_segments_fn = vmap(compute_planar_stiffness_matrix) - @jit def apply_eps_to_bend_strains(xi: Array, _eps: float) -> Array: """ Add a small number to the bending strain to avoid singularities """ - xi_reshaped = xi.reshape((-1, 3)) - - xi_bend_sign = jnp.sign(xi_reshaped[:, 0]) - # set zero sign to 1 (i.e. positive) - xi_bend_sign = jnp.where(xi_bend_sign == 0, 1, xi_bend_sign) - # add eps to the bending strain (i.e. the first column) - sigma_b_epsed = lax.select( - jnp.abs(xi_reshaped[:, 0]) < _eps, - xi_bend_sign * _eps, - xi_reshaped[:, 0], - ) - xi_epsed = jnp.stack( - [ - sigma_b_epsed, - xi_reshaped[:, 1], - xi_reshaped[:, 2], - ], - axis=1, - ) - - # old implementation: - # xi_epsed = xi_reshaped - # xi_epsed = xi_epsed.at[:, 0].add(xi_bend_sign * _eps) - - # flatten the array - xi_epsed = xi_epsed.flatten() + if _eps == None: + return xi + else: + xi_reshaped = xi.reshape((-1, 3)) + + xi_bend_sign = jnp.sign(xi_reshaped[:, 0]) + # set zero sign to 1 (i.e. positive) + xi_bend_sign = jnp.where(xi_bend_sign == 0, 1, xi_bend_sign) + # add eps to the bending strain (i.e. the first column) + sigma_b_epsed = lax.select( + jnp.abs(xi_reshaped[:, 0]) < _eps, + xi_bend_sign * _eps, + xi_reshaped[:, 0], + ) + xi_epsed = jnp.stack( + [ + sigma_b_epsed, + xi_reshaped[:, 1], + xi_reshaped[:, 2], + ], + axis=1, + ) + + # old implementation: + # xi_epsed = xi_reshaped + # xi_epsed = xi_epsed.at[:, 0].add(xi_bend_sign * _eps) + + # flatten the array + xi_epsed = xi_epsed.flatten() return xi_epsed @@ -206,8 +205,11 @@ def classify_segment(params: Dict[str, Array], s: Array) -> Tuple[Array, Array]: return segment_idx, s_segment if stiffness_fn is None: + def stiffness_fn( - params: Dict[str, Array], B_xi: Array, formulate_in_strain_space: bool = False + params: Dict[str, Array], + B_xi: Array, + formulate_in_strain_space: bool = False, ) -> Array: """ Compute the stiffness matrix of the system. @@ -237,6 +239,7 @@ def stiffness_fn( return S if actuation_mapping_fn is None: + def actuation_mapping_fn( forward_kinematics_fn: Callable, jacobian_fn: Callable, @@ -262,7 +265,6 @@ def actuation_mapping_fn( return A - @jit def forward_kinematics_fn( params: Dict[str, Array], q: Array, s: Array, eps: float = global_eps ) -> Array: @@ -296,7 +298,6 @@ def forward_kinematics_fn( return chi - @jit def jacobian_fn( params: Dict[str, Array], q: Array, s: Array, eps: float = global_eps ) -> Array: @@ -332,7 +333,6 @@ def jacobian_fn( return J - @jit def dynamical_matrices_fn( params: Dict[str, Array], q: Array, q_d: Array, eps: float = 1e4 * global_eps ) -> Tuple[Array, Array, Array, Array, Array, Array]: @@ -361,7 +361,9 @@ def dynamical_matrices_fn( # compute the stiffness matrix K = stiffness_fn(params, B_xi, formulate_in_strain_space=True) # compute the actuation matrix - A = actuation_mapping_fn(forward_kinematics_fn, jacobian_fn, params, B_xi, xi_eq, q) + A = actuation_mapping_fn( + forward_kinematics_fn, jacobian_fn, params, B_xi, xi_eq, q + ) # dissipative matrix from the parameters D = params.get("D", jnp.zeros((n_xi, n_xi))) @@ -474,7 +476,7 @@ def operational_space_dynamical_matrices_fn( Lambda: inertia matrix in the operational space of shape (3, 3) mu: matrix with corioli and centrifugal terms in the operational space of shape (3, 3) J: Jacobian of the Cartesian pose with respect to the generalized coordinates of shape (3, n_q) - J: time-derivative of the Jacobian of the end-effector pose with respect to the generalized coordinates + J_d: time-derivative of the Jacobian of the end-effector pose with respect to the generalized coordinates of shape (3, n_q) JB_pinv: Dynamically-consistent pseudo-inverse of the Jacobian. Allows the mapping of torques from the generalized coordinates to the operational space: f = JB_pinv.T @ tau_q diff --git a/src/jsrm/systems/pneumatically_actuated_planar_pcs.py b/src/jsrm/systems/pneumatically_actuated_planar_pcs.py index 0c91d5b..4e867f5 100644 --- a/src/jsrm/systems/pneumatically_actuated_planar_pcs.py +++ b/src/jsrm/systems/pneumatically_actuated_planar_pcs.py @@ -5,14 +5,15 @@ import numpy as onp from typing import Callable, Dict, Optional, Tuple, Union -from .planar_pcs import factory as planar_pcs_factory +from .planar_pcs_sym import factory as planar_pcs_factory + def factory( num_segments: int, *args, segment_actuation_selector: Optional[Array] = None, simplified_actuation_mapping: bool = False, - **kwargs + **kwargs, ): """ Factory function for the pneumatically-actuated planar PCS. @@ -69,9 +70,13 @@ def actuation_mapping_fn( J_sms = vmap(jacobian_fn, in_axes=(None, None, 0))(params, q, sms) def compute_actuation_matrix_for_segment( - r_cham_in: Array, r_cham_out: Array, varphi_cham: Array, - chi_pe: Array, chi_de: Array, - J_pe: Array, J_de: Array, + r_cham_in: Array, + r_cham_out: Array, + varphi_cham: Array, + chi_pe: Array, + chi_de: Array, + J_pe: Array, + J_de: Array, ) -> Array: """ Compute the actuation matrix for a single segment. @@ -100,32 +105,42 @@ def compute_actuation_matrix_for_segment( th_pe, th_de = chi_pe[2], chi_de[2] # compute the area of each pneumatic chamber (we assume identical chambers within a segment) - A_cham = 0.5 * varphi_cham * (r_cham_out ** 2 - r_cham_in ** 2) + A_cham = 0.5 * varphi_cham * (r_cham_out**2 - r_cham_in**2) # compute the center of pressure of the pneumatic chamber r_cop = ( - 2 / 3 * jnp.sinc(0.5 * varphi_cham) * (r_cham_out ** 3 - r_cham_in ** 3) / (r_cham_out ** 2 - r_cham_in ** 2) + 2 + / 3 + * jnp.sinc(0.5 * varphi_cham) + * (r_cham_out**3 - r_cham_in**3) + / (r_cham_out**2 - r_cham_in**2) ) if simplified_actuation_mapping: - A_sm = B_xi.T @ jnp.array([ - [A_cham * r_cop, -A_cham * r_cop], - [0.0, 0.0], - [2 * A_cham, 2 * A_cham], - ]) + A_sm = B_xi.T @ jnp.array( + [ + [A_cham * r_cop, -A_cham * r_cop], + [0.0, 0.0], + [2 * A_cham, 2 * A_cham], + ] + ) else: # compute the actuation matrix that collects the contributions of the pneumatic chambers in the given segment # first we consider the contribution of the distal end - A_sm_de = J_de.T @ jnp.array([ - [-2 * A_cham * jnp.sin(th_de), -2 * A_cham * jnp.sin(th_de)], - [2 * A_cham * jnp.cos(th_de), 2 * A_cham * jnp.cos(th_de)], - [A_cham * r_cop, -A_cham * r_cop] - ]) + A_sm_de = J_de.T @ jnp.array( + [ + [-2 * A_cham * jnp.sin(th_de), -2 * A_cham * jnp.sin(th_de)], + [2 * A_cham * jnp.cos(th_de), 2 * A_cham * jnp.cos(th_de)], + [A_cham * r_cop, -A_cham * r_cop], + ] + ) # then, we consider the contribution of the proximal end - A_sm_pe = J_pe.T @ jnp.array([ - [2 * A_cham * jnp.sin(th_pe), 2 * A_cham * jnp.sin(th_pe)], - [-2 * A_cham * jnp.cos(th_pe), -2 * A_cham * jnp.cos(th_pe)], - [-A_cham * r_cop, A_cham * r_cop] - ]) + A_sm_pe = J_pe.T @ jnp.array( + [ + [2 * A_cham * jnp.sin(th_pe), 2 * A_cham * jnp.sin(th_pe)], + [-2 * A_cham * jnp.cos(th_pe), -2 * A_cham * jnp.cos(th_pe)], + [-A_cham * r_cop, A_cham * r_cop], + ] + ) # sum the contributions of the distal and proximal ends A_sm = A_sm_de + A_sm_pe @@ -133,9 +148,13 @@ def compute_actuation_matrix_for_segment( return A_sm A_sms = vmap(compute_actuation_matrix_for_segment)( - params["r_cham_in"], params["r_cham_out"], params["varphi_cham"], - chi_pe=chi_sms[:-1], chi_de=chi_sms[1:], - J_pe=J_sms[:-1], J_de=J_sms[1:], + params["r_cham_in"], + params["r_cham_out"], + params["varphi_cham"], + chi_pe=chi_sms[:-1], + chi_de=chi_sms[1:], + J_pe=J_sms[:-1], + J_de=J_sms[1:], ) # we need to sum the contributions of the actuation of each segment A = jnp.sum(A_sms, axis=0) @@ -146,16 +165,31 @@ def compute_actuation_matrix_for_segment( return A return planar_pcs_factory( - *args, stiffness_fn=stiffness_fn, actuation_mapping_fn=actuation_mapping_fn, **kwargs + *args, + stiffness_fn=stiffness_fn, + actuation_mapping_fn=actuation_mapping_fn, + **kwargs, ) + def _compute_stiffness_matrix_for_segment( - l: Array, r: Array, r_cham_in: Array, r_cham_out: Array, varphi_cham: Array, E: Array + l: Array, + r: Array, + r_cham_in: Array, + r_cham_out: Array, + varphi_cham: Array, + E: Array, ): # cross-sectional area [m²] of the material - A_mat = jnp.pi * r ** 2 + 2 * r_cham_in ** 2 * varphi_cham - 2 * r_cham_out ** 2 * varphi_cham + A_mat = ( + jnp.pi * r**2 + 2 * r_cham_in**2 * varphi_cham - 2 * r_cham_out**2 * varphi_cham + ) # second moment of area [m⁴] of the material - Ib_mat = jnp.pi * r ** 4 / 4 + r_cham_in ** 4 * varphi_cham / 2 - r_cham_out ** 4 * varphi_cham / 2 + Ib_mat = ( + jnp.pi * r**4 / 4 + + r_cham_in**4 * varphi_cham / 2 + - r_cham_out**4 * varphi_cham / 2 + ) # poisson ratio of the material nu = 0.0 # shear modulus @@ -172,6 +206,7 @@ def _compute_stiffness_matrix_for_segment( return S + def stiffness_fn( params: Dict[str, Array], B_xi: Array, @@ -187,10 +222,13 @@ def stiffness_fn( S: elastic matrix of shape (n_q, n_q) if formulate_in_strain_space is False or (n_xi, n_xi) otherwise """ # stiffness matrix of shape (num_segments, 3, 3) - S_sms = vmap( - _compute_stiffness_matrix_for_segment - )( - params["l"], params["r"], params["r_cham_in"], params["r_cham_out"], params["varphi_cham"], params["E"] + S_sms = vmap(_compute_stiffness_matrix_for_segment)( + params["l"], + params["r"], + params["r_cham_in"], + params["r_cham_out"], + params["varphi_cham"], + params["E"], ) # we define the elastic matrix of shape (n_xi, n_xi) as K(xi) = S @ xi where K is equal to S = blk_diag(S_sms) @@ -198,4 +236,4 @@ def stiffness_fn( if not formulate_in_strain_space: S = B_xi.T @ S @ B_xi - return S \ No newline at end of file + return S diff --git a/src/jsrm/systems/tendon_actuated_planar_pcs.py b/src/jsrm/systems/tendon_actuated_planar_pcs.py index 28577f9..caee2df 100644 --- a/src/jsrm/systems/tendon_actuated_planar_pcs.py +++ b/src/jsrm/systems/tendon_actuated_planar_pcs.py @@ -5,13 +5,14 @@ import numpy as onp from typing import Callable, Dict, Optional, Tuple, Union -from .planar_pcs import factory as planar_pcs_factory +from .planar_pcs_sym import factory as planar_pcs_factory + def factory( num_segments: int, *args, segment_actuation_selector: Optional[Array] = None, - **kwargs + **kwargs, ): """ Factory function for the tendon-driven planar PCS. @@ -53,12 +54,13 @@ def actuation_mapping_fn( segment_indices = jnp.arange(num_segments) def compute_actuation_matrix_for_segment( - segment_idx, d_sm: Array, + segment_idx, + d_sm: Array, ) -> Array: """ Compute the actuation matrix for a single segment. - We assume that each segment is actuated by num_segment_tendons that are routed at a distance of d from the segment's backbone, - respectively, and attached to the segment's distal end. We assume that the motor is located at the base of the robot and that the + We assume that each segment is actuated by num_segment_tendons that are routed at a distance of d from the segment's backbone, + respectively, and attached to the segment's distal end. We assume that the motor is located at the base of the robot and that the tendons are routed through all proximal segments. The positive control inputs u1 and u2 are the tensions (i.e., forces) applied by the two tendons. At a straight configuration with a positive d1, a positive u1 and zero u2 should cause the bend negatively (to the right) and contract its length. @@ -69,6 +71,7 @@ def compute_actuation_matrix_for_segment( Returns: A_sm: actuation matrix of shape (n_xi, num_segment_tendons) """ + def compute_A_d(d: Array) -> Array: """ Compute the actuation matrix for a single actuator/tendon with respect to the soft robot's strains. @@ -77,6 +80,7 @@ def compute_A_d(d: Array) -> Array: Returns: A_d: actuation matrix of shape (n_xi, ) where n_xi is the number of strains """ + def compute_A_d_wrt_xi_i(i: Array, l_i: Array, xi_i: Array) -> Array: """ Compute the actuation matrix for a single actuator with respect to the strains of a single segment. @@ -87,22 +91,26 @@ def compute_A_d_wrt_xi_i(i: Array, l_i: Array, xi_i: Array) -> Array: Returns: A_d_segment: actuation matrix for the segment of shape (3, 3) """ - square_root_term = jnp.sqrt(xi[1]**2 + (xi[2] + d * xi[0])**2) - A_d_wrt_xi_i = - jnp.array([ - l_i * d * (d * xi_i[0] + xi_i[2]) / square_root_term, - l_i * xi_i[1] / square_root_term, - l_i * (d * xi_i[0] + xi_i[2]) / square_root_term, - ]) + square_root_term = jnp.sqrt(xi[1] ** 2 + (xi[2] + d * xi[0]) ** 2) + A_d_wrt_xi_i = -jnp.array( + [ + l_i * d * (d * xi_i[0] + xi_i[2]) / square_root_term, + l_i * xi_i[1] / square_root_term, + l_i * (d * xi_i[0] + xi_i[2]) / square_root_term, + ] + ) return jnp.where( - i * jnp.ones((3, )) <= segment_idx * jnp.ones((3, )), + i * jnp.ones((3,)) <= segment_idx * jnp.ones((3,)), A_d_wrt_xi_i, - jnp.zeros_like(A_d_wrt_xi_i) + jnp.zeros_like(A_d_wrt_xi_i), ) - A_d = vmap(compute_A_d_wrt_xi_i)(segment_indices, l, xi.reshape(-1, 3)).reshape(-1) + A_d = vmap(compute_A_d_wrt_xi_i)( + segment_indices, l, xi.reshape(-1, 3) + ).reshape(-1) return A_d - + A_sm = vmap(compute_A_d, in_axes=0, out_axes=-1)(d_sm) return A_sm @@ -110,9 +118,10 @@ def compute_A_d_wrt_xi_i(i: Array, l_i: Array, xi_i: Array) -> Array: # compute the actuation matrix for all segments # will have shape (num_segments, n_xi, num_segment_tendons) A = vmap(compute_actuation_matrix_for_segment, in_axes=(0, 0), out_axes=0)( - segment_indices, params["d"], + segment_indices, + params["d"], ) - + # deactivate the actuation for some segments A = A[segment_actuation_selector] diff --git a/src/jsrm/systems/utils.py b/src/jsrm/systems/utils.py index f1850ae..cd6d42b 100644 --- a/src/jsrm/systems/utils.py +++ b/src/jsrm/systems/utils.py @@ -1,13 +1,16 @@ from copy import deepcopy import jax -from jax import Array, jit, vmap + from jax import numpy as jnp import sympy as sp -from typing import Callable, Dict, Iterable, List, Sequence, Tuple, Union + +# For documentation purposes +from jax import Array +from typing import Dict, List, Tuple, Union def substitute_params_into_all_symbolic_expressions( - sym_exps: Dict, params: Dict[str, jnp.array] + sym_exps: Dict, params: Dict[str, Array] ) -> Dict: """ Substitute robot parameters into symbolic expressions. @@ -47,7 +50,7 @@ def substitute_params_into_all_symbolic_expressions( def substitute_params_into_single_symbolic_expression( sym_exp: sp.Expr, params_syms: Dict[str, List[sp.Symbol]], - params: Dict[str, jnp.array], + params: Dict[str, Array], ) -> sp.Expr: """ Substitute robot parameters into a single symbolic expression. @@ -88,13 +91,15 @@ def concatenate_params_syms( def compute_strain_basis( strain_selector: Array, -) -> jnp.ndarray: +) -> Array: """ Compute strain basis based on boolean strain selector. Args: - strain_selector: boolean array of shape (n_xi, ) specifying which strain components are active + strain_selector (Array): + boolean array of shape (n_xi, ) specifying which strain components are active Returns: - strain_basis: strain basis matrix of shape (n_xi, n_q) where n_q is the number of configuration variables + strain_basis (Array): + strain basis matrix of shape (n_xi, n_q) where n_q is the number of configuration variables and n_xi is the number of strains """ n_q = strain_selector.sum().item() @@ -107,8 +112,9 @@ def compute_strain_basis( return strain_basis -@jit -def compute_planar_stiffness_matrix(l: Array, A: Array, Ib: Array, E: Array, G: Array) -> Array: +def compute_planar_stiffness_matrix( + l: Array, A: Array, Ib: Array, E: Array, G: Array +) -> Array: """ Compute the stiffness matrix of the system. Args: @@ -124,3 +130,75 @@ def compute_planar_stiffness_matrix(l: Array, A: Array, Ib: Array, E: Array, G: S = l * jnp.diag(jnp.stack([Ib * E, 4 / 3 * A * G, A * E], axis=0)) return S + + +def gauss_quadrature(N_GQ: int, a=0.0, b=1.0) -> Tuple[Array, Array, int]: + """ + Computes the Legendre-Gauss nodes and weights on the interval [0, 1] + using Legendre-Gauss Quadrature with truncation order N_GQ. + + Args: + N_GQ (int): order of the truncature. + a (float, optional): The lower bound of the interval. Default is 0.0. + b (float, optional): The upper bound of the interval. Default is 1.0. + + Returns: + Xs (Array): The Gauss nodes on [a, b]. + Ws (Array): The Gauss weights on [a, b]. + nGauss (int): The number of Gauss points including boundary points, i.e., N_GQ + 2. + """ + + N = N_GQ - 1 + N1 = N + 1 + N2 = N + 2 + + xu = jnp.linspace(-1, 1, N1) + + # Initial guess + y = jnp.cos((2 * jnp.arange(N + 1) + 1) * jnp.pi / (2 * N + 2)) + ( + 0.27 / N1 + ) * jnp.sin(jnp.pi * xu * N / N2) + + def legendre_iteration(y): + L = [jnp.ones_like(y), y] + for k in range(2, N1 + 1): + Lk = ((2 * k - 1) * y * L[-1] - (k - 1) * L[-2]) / k + L.append(Lk) + L = jnp.stack(L, axis=1) + Lp = N2 * (L[:, N1 - 1] - y * L[:, N1]) / (1 - y**2) + return y - L[:, N1] / Lp + + def convergence_condition(y): + L = [jnp.ones_like(y), y] + for k in range(2, N1 + 1): + Lk = ((2 * k - 1) * y * L[-1] - (k - 1) * L[-2]) / k + L.append(Lk) + L = jnp.stack(L, axis=1) + Lp = N2 * (L[:, N1 - 1] - y * L[:, N1]) / (1 - y**2) + y_new = y - L[:, N1] / Lp + return jnp.max(jnp.abs(y_new - y)) > jnp.finfo(jnp.float32).eps + + y = jax.lax.while_loop( # TODO + convergence_condition, legendre_iteration, y + ) + + # Linear map from [-1, 1] to [a, b] + Xs = (a * (1 - y) + b * (1 + y)) / 2 + Xs = jnp.flip(Xs) + + # Add the boundary points + Xs = jnp.concatenate([jnp.array([a]), Xs, jnp.array([b])]) + + # Compute the weights + L = [jnp.ones_like(y), y] + for k in range(2, N1 + 1): + Lk = ((2 * k - 1) * y * L[-1] - (k - 1) * L[-2]) / k + L.append(Lk) + L = jnp.stack(L, axis=1) + Lp = N2 * (L[:, N1 - 1] - y * L[:, N1]) / (1 - y**2) + Ws = (b - a) / ((1 - y**2) * Lp**2) * (N2 / N1) ** 2 + + # Add the boundary points + Ws = jnp.concatenate([jnp.array([0.0]), Ws, jnp.array([0.0])]) + + return Xs, Ws, N_GQ + 2 diff --git a/src/jsrm/utils/lie_operators.py b/src/jsrm/utils/lie_operators.py new file mode 100644 index 0000000..4ca7dca --- /dev/null +++ b/src/jsrm/utils/lie_operators.py @@ -0,0 +1,608 @@ +import jax.numpy as jnp + +# for documentation +from jax import Array +from typing import Sequence + + +def tilde_SE3(vec3: Array) -> Array: + """ + Computes the tilde operator of SE(3) for a 3D vector. + + Args: + vec3 (Array): array-like, shape (3,1) + A 3-dimensional vector. + + Returns: + Array: shape (3, 3) + A 3x3 matrix representing the tilde operator of the input vector. + """ + vec3 = vec3.reshape(-1) # Ensure vec3 is a 1D array + + # Extract components of the vector + x, y, z = vec3.flatten() + + # Use JAX's array creation for better performance + Mtilde = jnp.array([[0, -z, y], [z, 0, -x], [-y, x, 0]]) + return Mtilde + + +def adjoint_SE3(vec6: Array) -> Array: + """ + Computes the adjoint representation of a vector of se(3). + + Args: + vec6 (Array): array-like, shape (6,1) + A 6-dimensional vector representing the screw. + The first three elements correspond to the angular component, + and the last three elements correspond to the linear component. + + Returns: + Array: shape (6, 6) + A 6x6 matrix representing the adjoint transformation of the input screw vector. + """ + vec6 = vec6.reshape(-1) # Ensure vec6 is a 1D array + + ang = vec6[:3].reshape((3, 1)) # Angular as a (3,1) vector + lin = vec6[3:].reshape((3, 1)) # Linear as a (3,1) vector + + angtilde = tilde_SE3(ang) # Tilde operator for angular part + lintilde = tilde_SE3(lin) # Tilde operator for linear part + + adj = jnp.block([[angtilde, jnp.zeros((3, 3))], [lintilde, angtilde]]) + + return adj + + +def adjoint_star_SE3(vec6: Array) -> Array: + """ + Computes the co-adjoint representation of a vector of se(3). + + Args: + vec6 (Array): array-like, shape (6,1) + A 6-dimensional vector representing the screw. + The first three elements correspond to the angular component, + and the last three elements correspond to the linear component. + + Returns: + Array: shape (6, 6) + A 6x6 matrix representing the co-adjoint transformation of the input screw vector. + """ + vec6 = vec6.reshape(-1) # Ensure vec6 is a 1D array + + ang = vec6[:3].reshape((3, 1)) # Angular as a (3,1) vector + lin = vec6[3:].reshape((3, 1)) # Linear as a (3,1) vector + + angtilde = tilde_SE3(ang) # Tilde operator for angular part + lintilde = tilde_SE3(lin) # Tilde operator for linear part + + adj_star = jnp.block([[angtilde, lintilde], [jnp.zeros((3, 3)), angtilde]]) + + return adj_star + + +def hat_SE3(vec6: Array) -> Array: + """ + Computes the hat operator for a 6D vector of se(3). + + Args: + vec6 (Array): array-like, shape (6,1) + A 6-dimensional vector representing the screw. + The first three elements correspond to the angular component, + and the last three elements correspond to the linear component. + + Returns: + Array: shape (4, 4) + A 4x4 matrix representing the hat operator of the input screw vector. + """ + vec6 = vec6.reshape(-1) # Ensure vec6 is a 1D array + + ang = vec6[:3].reshape((3, 1)) # Angular as a (3,1) vector + lin = vec6[3:].reshape((3, 1)) # Linear as a (3,1) vector + + angtilde = tilde_SE3(ang) # Tilde operator for angular part + + hat = jnp.block([[angtilde, lin], [jnp.zeros((1, 3)), jnp.zeros((1, 1))]]) + + return hat + + +def Adjoint_g_SE3(mat4: Array) -> Array: + """ + Computes the adjoint representation of a 4x4 matrix. + + Args: + mat4 (Array): array-like, shape (4,4) + A 4x4 matrix representing the transformation. + + Returns: + Array: shape (4, 4) + A 4x4 matrix representing the Adjoint transformation of the input matrix. + """ + ang = mat4[:3, :3] # Extract the angular part (top-left 3x3 block) + lin = mat4[:3, 3].reshape((3, 1)) # Extract the linear part (top-right column) + + ltilde = tilde_SE3(lin) # Tilde operator for linear part + + Adjoint = jnp.block([[ang, jnp.zeros((3, 3))], [ltilde @ ang, ang]]) + + return Adjoint + + +def Adjoint_gn_SE3( + xi_n: Array, + l_nprev: float, + s: float, +) -> Array: + """ + Computes the adjoint representation of a position of a points at s (general curvilinear coordinate) + along a rod in SE(3) deformed ine the current segment according to a strain vector xi_n. + + If s is a point of the n-th segment, this function use the lenght from the origin of the rod to the begining of the n-th segment. + + Args: + xi_n (Array): array-like, shape (6,1) + A 6-dimensional vector representing the screw in SE(3). + l_nprev (float): + The length from the origin of the rod to the beginning of the n-th segment. + s (float): + The curvilinear coordinate along the rod, representing the position of a point in the n-th segment. + + Returns: + Array: shape (6, 6) + A 6x6 matrix representing the adjoint transformation of the input screw vector at the specified position. + """ + # We suppose here that theta is not zero thanks to a previous use of apply_eps + ang = xi_n[:3].reshape((3, 1)) # Angular as a (3,1) vector + theta = jnp.linalg.norm(ang) # Compute the norm of the angular part + x = s - l_nprev # Compute the segment length + adjoint_xi_n = adjoint_SE3(xi_n) # Adjoint representation of the input vector + + cos = jnp.cos(x * theta) + sin = jnp.sin(x * theta) + Adjoint = ( + jnp.eye(6) + + 1 / (2 * theta) * (3 * sin - x * theta * cos) * adjoint_xi_n + + 1 + / (2 * jnp.power(theta, 2)) + * (4 - 4 * cos - x * theta * sin) + * jnp.linalg.matrix_power(adjoint_xi_n, 2) + + 1 + / (2 * jnp.power(theta, 3)) + * (sin - x * theta * cos) + * jnp.linalg.matrix_power(adjoint_xi_n, 3) + + 1 + / (2 * jnp.power(theta, 4)) + * (2 - 2 * cos - x * theta * sin) + * jnp.linalg.matrix_power(adjoint_xi_n, 4) + ) + + return Adjoint + + +def Adjoint_gn_SE3_inv( + xi_n: Array, + l_nprev: float, + s: float, +) -> Array: + """ + Computes the adjoint representation of a position of a points at s (general curvilinear coordinate) + along a rod in SE(3) deformed ine the current segment according to a strain vector xi_n. + + If s is a point of the n-th segment, this function use the lenght from the origin of the rod to the begining of the n-th segment. + + Args: + xi_n (Array): array-like, shape (6,1) + A 6-dimensional vector representing the screw in SE(3). + l_nprev (float): + The length from the origin of the rod to the beginning of the n-th segment. + s (float): + The curvilinear coordinate along the rod, representing the position of a point in the n-th segment. + + Returns: + Array: shape (6, 6) + A 6x6 matrix representing the adjoint transformation of the input screw vector at the specified position. + """ + # We suppose here that theta is not zero thanks to a previous use of apply_eps + ang = xi_n[:3].reshape((3, 1)) # Angular as a (3,1) vector + theta = jnp.linalg.norm(ang) # Compute the norm of the angular part + x = s - l_nprev # Compute the segment length + adjoint_xi_n = adjoint_SE3(xi_n) # Adjoint representation of the input vector + + cos = jnp.cos(x * theta) + sin = jnp.sin(x * theta) + + Adjoint = ( + jnp.eye(6) + + 1 / (2 * theta) * (3 * sin - x * theta * cos) * adjoint_xi_n + + 1 + / (2 * jnp.power(theta, 2)) + * (4 - 4 * cos - x * theta * sin) + * jnp.linalg.matrix_power(adjoint_xi_n, 2) + + 1 + / (2 * jnp.power(theta, 3)) + * (sin - x * theta * cos) + * jnp.linalg.matrix_power(adjoint_xi_n, 3) + + 1 + / (2 * jnp.power(theta, 4)) + * (2 - 2 * cos - x * theta * sin) + * jnp.linalg.matrix_power(adjoint_xi_n, 4) + ) + + # Extract R and uR from the Adjoint matrix + R = Adjoint[:3, :3] + uR = Adjoint[3:, :3] + + # Compute the inverse using the Schur complement + R_inv = jnp.transpose(R) # Since R is a rotation matrix + u = jnp.dot(uR, R_inv) # Compute the linear part + uR_inv = -jnp.dot(R_inv, u) # Compute the inverse linear part + + # Construct the inverse Adjoint matrix + inverse_Adjoint = jnp.block([[R_inv, jnp.zeros((3, 3))], [uR_inv, R_inv]]) + + return inverse_Adjoint + + +def Tangent_gn_SE3( + xi_n: Array, + l_nprev: float, + s: float, +) -> Array: + """ + Computes the tangent representation of a position of a points at s (general curvilinear coordinate) + along a rod in SE(3) deformed in the current segment according to a strain vector xi_n. + + If s is a point of the n-th segment, this function use the length from the origin of the rod to the beginning of the n-th segment. + + Args: + xi_n (Array): array-like, shape (6,1) + A 6-dimensional vector representing the screw in SE(3). + l_nprev (float): + The length from the origin of the rod to the beginning of the n-th segment. + s (float): + The curvilinear coordinate along the rod, representing the position of a point in the n-th segment. + + Returns: + Array: shape (6, 6) + A 6x6 matrix representing the tangent transformation of the input screw vector at the specified position. + """ + # We suppose here that theta is not zero thanks to a previous use of apply_eps + ang = xi_n[:3].reshape((3, 1)) # Angular as a (3,1) vector + theta = jnp.linalg.norm(ang) # Compute the norm of the angular part + x = s - l_nprev # Compute the segment length + adjoint_xi_n = adjoint_SE3(xi_n) # Adjoint representation of the input vector + + cos = jnp.cos(x * theta) + sin = jnp.sin(x * theta) + + Tangent = ( + x * jnp.eye(6) + + 1 / (2 * jnp.power(theta, 2)) * (4 - 4 * cos - x * theta * sin) * adjoint_xi_n + + 1 + / (2 * jnp.power(theta, 3)) + * (4 * x * theta - 5 * sin + x * theta * cos) + * jnp.linalg.matrix_power(adjoint_xi_n, 2) + + 1 + / (2 * jnp.power(theta, 4)) + * (2 - 2 * cos - x * theta * sin) + * jnp.linalg.matrix_power(adjoint_xi_n, 3) + + 1 + / (2 * jnp.power(theta, 5)) + * (2 * x * theta - 3 * sin + x * theta * cos) + * jnp.linalg.matrix_power(adjoint_xi_n, 4) + ) + + return Tangent + + +def vec_SE2_to_xi_SE3(vec3: Array, indices: Sequence[int] = (2, 3, 4)) -> Array: + """ + Convert a strain vector in se(2) to a strain vector in se(3). + + Args: + vec3 (Array): array-like, shape (3,1) + A 3-dimensional vector representing the strain in se(2). + The first element correspond to the angular component, + and the last elements corresponds to the linear component. + indices (Sequence[int], optional): Indices in the 6D se(3) vector + where to insert the se(2) components. Default is (2, 3, 4) + + Returns: + Array: shape (6,1) + A 6-dimensional vector representing the strain in se(3). + The first three elements correspond to the angular component, + and the last three elements correspond to the linear component. + """ + vec3 = jnp.asarray(vec3).flatten() # Ensure vec3 is a JAX array + + xi = jnp.zeros((6,)) # Initialize a 6D vector with zeros + xi = xi.at[jnp.array(indices)].set(vec3) # Set the values at the specified indices + return xi.reshape((6, 1)) + + +# ================================================================================================ +# SE(2) operators +# =================================== +J = jnp.array([[0, -1], [1, 0]]) + + +def adjoint_SE2(vec3: Array) -> Array: + """ + Computes the adjoint representation of a vector of se(2). + + Args: + vec3 (Array): array-like, shape (3, 1) + A 3-dimensional vector representing the screw. + The first element correspond to the angular component, + and the last two elements correspond to the linear component. + + Returns: + Array: shape (3, 3) + A 3x3 matrix representing the adjoint transformation of the input screw vector. + """ + vec3 = vec3.reshape(-1) # Ensure vec6 is a 1D array + + ang = vec3[0] + lin = vec3[1:].reshape((2, 1)) # Linear as a (3,1) vector + + adj = jnp.concatenate( + [jnp.zeros((1, 3)), jnp.concatenate([-J @ lin, ang * J], axis=1)] + ) + + return adj + + +def adjoint_star_SE2(vec3: Array) -> Array: + """ + Computes the co-adjoint representation of a vector of se(2). + + Args: + vec3 (Array): array-like, shape (3, 1) + A 3-dimensional vector representing the screw. + The first element correspond to the angular component, + and the last two elements correspond to the linear component. + + Returns: + Array: shape (3, 3) + A 3x3 matrix representing the co-adjoint transformation of the input screw vector. + """ + vec3 = vec3.reshape(-1) # Ensure vec6 is a 1D array + + ang = vec3[0] + lin = vec3[1:].reshape((2, 1)) # Linear as a (3,1) vector + + adj_star = jnp.concatenate( + [jnp.zeros((3, 1)), jnp.concatenate([lin.T @ J, ang * J], axis=0)], axis=1 + ) + + return adj_star + + +def Adjoint_g_SE2(mat3: Array) -> Array: + """ + Computes the adjoint representation of a 3x3 matrix. + + Args: + mat4 (Array): array-like, shape (4,4) + A 4x4 matrix representing the transformation. + + Returns: + Array: shape (3, 3) + A 3x3 matrix representing the Adjoint transformation of the input matrix. + """ + R = mat3[:2, :2] # Extract the angular part (top-left 2x2 block) + t = mat3[:2, 2].reshape((2, 1)) # Extract the linear part (top-right column) + + Adjoint = jnp.concatenate( + [ + jnp.concatenate([jnp.ones(((1, 1))), jnp.zeros((1, 2))], axis=1), + jnp.concatenate([-J @ t, R], axis=1), + ] + ) + + return Adjoint + + +def Adjoint_gn_SE2( + xi_n: Array, + l_nprev: float, + s: float, +) -> Array: + """ + Computes the adjoint representation of a position of a points at s (general curvilinear coordinate) + along a rod in SE(2) deformed ine the current segment according to a strain vector xi_n. + + If s is a point of the n-th segment, this function use the lenght from the origin of the rod to the begining of the n-th segment. + + Args: + xi_n (Array): array-like, shape (3,1) + A 3-dimensional vector representing the screw in SE(2). + l_nprev (float): + The length from the origin of the rod to the beginning of the n-th segment. + s (float): + The curvilinear coordinate along the rod, representing the position of a point in the n-th segment. + + Returns: + Array: shape (3, 3) + A 3x3 matrix representing the adjoint transformation of the input screw vector at the specified position. + """ + # We suppose here that theta is not zero thanks to a previous use of apply_eps + theta = xi_n[0] # Angular part + x = s - l_nprev # Compute the segment length + adjoint_xi_n = adjoint_SE2(xi_n) # Adjoint representation of the input vector + + cos = jnp.cos(x * theta) + sin = jnp.sin(x * theta) + + Adjoint = ( + jnp.eye(3) + + 1 / (2 * theta) * (3 * sin - x * theta * cos) * adjoint_xi_n + + 1 + / (2 * jnp.power(theta, 2)) + * (4 - 4 * cos - x * theta * sin) + * jnp.linalg.matrix_power(adjoint_xi_n, 2) + + 1 + / (2 * jnp.power(theta, 3)) + * (sin - x * theta * cos) + * jnp.linalg.matrix_power(adjoint_xi_n, 3) + + 1 + / (2 * jnp.power(theta, 4)) + * (2 - 2 * cos - x * theta * sin) + * jnp.linalg.matrix_power(adjoint_xi_n, 4) + ) + + return Adjoint + + +def Adjoint_gn_SE2_inv( + xi_n: Array, + l_nprev: float, + s: float, +) -> Array: + """ + Computes the adjoint representation of a position of a points at s (general curvilinear coordinate) + along a rod in SE(2) deformed ine the current segment according to a strain vector xi_n. + + If s is a point of the n-th segment, this function use the lenght from the origin of the rod to the begining of the n-th segment. + + Args: + xi_n (Array): array-like, shape (3,1) + A 3-dimensional vector representing the screw in SE(2). + l_nprev (float): + The length from the origin of the rod to the beginning of the n-th segment. + s (float): + The curvilinear coordinate along the rod, representing the position of a point in the n-th segment. + + Returns: + Array: shape (3, 3) + A 3x3 matrix representing the adjoint transformation of the input screw vector at the specified position. + """ + # We suppose here that theta is not zero thanks to a previous use of apply_eps + theta = xi_n[0] # Angular part + x = s - l_nprev # Compute the segment length + adjoint_xi_n = adjoint_SE2(xi_n) # Adjoint representation of the input vector + + cos = jnp.cos(x * theta) + sin = jnp.sin(x * theta) + + Adjoint = ( + jnp.eye(3) + + 1 / (2 * theta) * (3 * sin - x * theta * cos) * adjoint_xi_n + + 1 + / (2 * jnp.power(theta, 2)) + * (4 - 4 * cos - x * theta * sin) + * jnp.linalg.matrix_power(adjoint_xi_n, 2) + + 1 + / (2 * jnp.power(theta, 3)) + * (sin - x * theta * cos) + * jnp.linalg.matrix_power(adjoint_xi_n, 3) + + 1 + / (2 * jnp.power(theta, 4)) + * (2 - 2 * cos - x * theta * sin) + * jnp.linalg.matrix_power(adjoint_xi_n, 4) + ) + + # Extract R and -Jt from the Adjoint matrix + R = Adjoint[1:, 1:] + mJt = Adjoint[1:, 0].reshape(-1, 1) + + # Compute the inverse using the Schur complement + R_inv = jnp.transpose(R) # Since R is a rotation matrix, R^-1=R^T + # Construct the inverse Adjoint matrix + inverse_Adjoint = jnp.concatenate( + [ + jnp.concatenate([jnp.ones(((1, 1))), jnp.zeros((1, 2))], axis=1), + jnp.concatenate([-R_inv @ mJt, R_inv], axis=1), + ] + ) + + return inverse_Adjoint + + +def Tangent_gn_SE2( + xi_n: Array, + l_nprev: float, + s: float, +) -> Array: + """ + Computes the tangent representation of a position of a points at s (general curvilinear coordinate) + along a rod in SE(2) deformed in the current segment according to a strain vector xi_n. + + If s is a point of the n-th segment, this function use the length from the origin of the rod to the beginning of the n-th segment. + + Args: + xi_n (Array): array-like, shape (3,1) + A 3-dimensional vector representing the screw in SE(2). + l_nprev (float): + The length from the origin of the rod to the beginning of the n-th segment. + s (float): + The curvilinear coordinate along the rod, representing the position of a point in the n-th segment. + + Returns: + Array: shape (3, 3) + A 3x3 matrix representing the tangent transformation of the input screw vector at the specified position. + """ + # We suppose here that theta is not zero thanks to a previous use of apply_eps + theta = xi_n[0] # Angular part + x = s - l_nprev # Compute the segment length + adjoint_xi_n = adjoint_SE2(xi_n) # Adjoint representation of the input vector + + cos = jnp.cos(x * theta) + sin = jnp.sin(x * theta) + + Tangent = ( + x * jnp.eye(3) + + 1 / (2 * jnp.power(theta, 2)) * (4 - 4 * cos - x * theta * sin) * adjoint_xi_n + + 1 + / (2 * jnp.power(theta, 3)) + * (4 * x * theta - 5 * sin + x * theta * cos) + * jnp.linalg.matrix_power(adjoint_xi_n, 2) + + 1 + / (2 * jnp.power(theta, 4)) + * (2 - 2 * cos - x * theta * sin) + * jnp.linalg.matrix_power(adjoint_xi_n, 3) + + 1 + / (2 * jnp.power(theta, 5)) + * (2 * x * theta - 3 * sin + x * theta * cos) + * jnp.linalg.matrix_power(adjoint_xi_n, 4) + ) + + return Tangent + + +# ================================================================================================ +# Shared operators +# ============================ +def compute_weighted_sums(M: Array, vecm: Array, idx: int) -> Array: + """ + Compute the weighted sums of the matrix product of M and vecm, + + Args: + M (Array): array of shape (N, m, m) + Describes the matrix to be multiplied with vecm + vecm (Array): array-like of shape (N, m) + Describes the vector to be multiplied with M + idx (int): index of the last row to be summed over + + Returns: + Array: array of shape (N, m) + The result of the weighted sums. For each i, the result is the sum of the products of M[i, j] and vecm[j] for j from 0 to idx. + """ + N = M.shape[0] + # Matrix product for each j: (N, m, m) @ (N, m, 1) -> (N, m) + prod = jnp.einsum("nij,nj->ni", M, vecm) + + # Triangular mask for partial sum: (N, N) + # mask[i, j] = 1 if j >= i and j <= idx + mask = (jnp.arange(N)[:, None] <= jnp.arange(N)[None, :]) & ( + jnp.arange(N)[None, :] <= idx + ) + mask = mask.astype(M.dtype) # (N, N) + + # Extend 6-dimensional mask (N, N, 1) to apply to (N, m) + masked_prod = mask[:, :, None] * prod[None, :, :] # (N, N, m) + + # Sum over j for each i : (N, m) + result = masked_prod.sum(axis=1) # (N, m) + return result diff --git a/src/jsrm/utils/numerical_jacobian.py b/src/jsrm/utils/numerical_jacobian.py index 56d34a9..13e6c4f 100644 --- a/src/jsrm/utils/numerical_jacobian.py +++ b/src/jsrm/utils/numerical_jacobian.py @@ -137,7 +137,7 @@ def _eps_for_method(x0_dtype, f0_dtype, method): return EPS ** (1 / 3) else: raise RuntimeError( - "Unknown step method, should be one of " "{'2-point', '3-point'}" + "Unknown step method, should be one of {'2-point', '3-point'}" ) diff --git a/tests/gui_test_fwd_kine_eps_planar_pcs.py b/tests/gui_test_fwd_kine_eps_planar_pcs.py new file mode 100644 index 0000000..26e0bec --- /dev/null +++ b/tests/gui_test_fwd_kine_eps_planar_pcs.py @@ -0,0 +1,250 @@ +import jax.numpy as jnp +import matplotlib.pyplot as plt +from matplotlib.patches import Patch +from matplotlib.widgets import Slider, Button, CheckButtons +from matplotlib import rc + +rc("animation", html="html5") + +from jsrm.systems import planar_pcs_num, planar_pcs_sym +from pathlib import Path +import jsrm +from tqdm import tqdm +from matplotlib.animation import FFMpegWriter + +# === Paramètres initiaux === +num_segments = 1 +rho = 1070 * jnp.ones((num_segments,)) +params = { + "th0": jnp.array(0.0), + "l": 1e-1 * jnp.ones((num_segments,)), + "r": 2e-2 * jnp.ones((num_segments,)), + "rho": rho, + "g": jnp.array([0.0, 9.81]), + "E": 2e3 * jnp.ones((num_segments,)), + "G": 1e3 * jnp.ones((num_segments,)), +} +params["D"] = 1e-3 * jnp.diag( + ( + jnp.repeat(jnp.array([[1e0, 1e3, 1e3]]), num_segments, axis=0) + * params["l"][:, None] + ).flatten() +) +strain_selector = jnp.ones((3 * num_segments,), dtype=bool) + + +# === Chargement des fonctions Jacobiennes === +def get_fwd_kine_fn(jacobian_type): + if jacobian_type == "symbolic": + sym_exp_filepath = ( + Path(jsrm.__file__).parent + / "symbolic_expressions" + / f"planar_pcs_ns-{num_segments}.dill" + ) + _, fwd, _, _ = planar_pcs_sym.factory(sym_exp_filepath, strain_selector) + else: + _, fwd, _, _ = planar_pcs_num.factory( + num_segments, + strain_selector, + integration_type="gauss-legendre", + param_integration=5, + jacobian_type=jacobian_type, + ) + return fwd + + +FwdKine_autodiff_fn = get_fwd_kine_fn("autodiff") +FwdKine_explicit_fn = get_fwd_kine_fn("explicit") +FwdKine_symbolic_fn = get_fwd_kine_fn("symbolic") + +jacobian_colors = {"symbolic": "green", "explicit": "orange", "autodiff": "blue"} +jacobian_markers = {"symbolic": "s", "explicit": "x", "autodiff": "o"} +jacobian_types = ["symbolic", "explicit", "autodiff"] +list_of_type_of_jacobian = jacobian_types.copy() + +# === eps discret === +eps_options = [None, 1e-6, 1e-5, 1e-4, 1e-3, 1e-2, 1e-1] +eps_labels = ["None"] + [f"1e-{i}" for i in range(6, 0, -1)] + + +def get_eps_from_slider(): + return eps_options[int(eps_slider.val)] + + +# === Variables physiques === +borne_kappa = 1e-3 +nb_kappa = 51 +kappa_values = jnp.linspace(-borne_kappa, borne_kappa, nb_kappa) + +borne_sigma_x, borne_sigma_y = 1e-1, 1e-1 +nb_sigma_x, nb_sigma_y, nb_s = 50, 50, 50 +sigma_x_values = jnp.linspace(0, borne_sigma_x, nb_sigma_x) +sigma_y_values = jnp.linspace(0, borne_sigma_y, nb_sigma_y) +s_values = jnp.linspace(0, float(params["l"][0]), nb_s) + +# === Tracé principal === +fig, axs = plt.subplots(2, figsize=(15, 8)) + + +def FwdKine_plot(eps_list, s, sigma_x, sigma_y, fig, axs): + for ax in axs: + ax.clear() + ax.ticklabel_format(style="sci", axis="both", scilimits=(0, 0)) + ax.set_xlabel("kappa (bending strain)") + ax.set_ylabel("Fwd Kinematics components") + + for i_eps, eps in enumerate(eps_list): + FwdKine_kwargs = {"eps": eps} if eps is not None else {} + FwdKine_auto, FwdKine_exp, FwdKine_symb = [], [], [] + for kappa in kappa_values: + q = jnp.array([kappa, sigma_x, sigma_y - 1.0] * num_segments) + FwdKine_auto.append(FwdKine_autodiff_fn(params, q, s, **FwdKine_kwargs)) + FwdKine_exp.append(FwdKine_explicit_fn(params, q, s, **FwdKine_kwargs)) + FwdKine_symb.append(FwdKine_symbolic_fn(params, q, s, **FwdKine_kwargs)) + FwdKine_auto, FwdKine_exp, FwdKine_symb = ( + jnp.stack(FwdKine_auto), + jnp.stack(FwdKine_exp), + jnp.stack(FwdKine_symb), + ) + + for i in range(2): + if eps is not None: + axs[i].axvline( + eps, + color="red", + linestyle=":", + linewidth=2, + alpha=(i_eps + 1) / len(eps_list), + label=f"+/-eps={eps:.2e}", + ) + axs[i].axvline( + -eps, + color="red", + linestyle=":", + linewidth=2, + alpha=(i_eps + 1) / len(eps_list), + ) + + for j_type in list_of_type_of_jacobian: + FwdKine = { + "symbolic": FwdKine_symb, + "explicit": FwdKine_exp, + "autodiff": FwdKine_auto, + }[j_type] + axs[i].plot( + kappa_values, + FwdKine[:, i], + marker=jacobian_markers[j_type], + label=f"FwdKine_{j_type}", + color=jacobian_colors[j_type], + alpha=(i_eps + 1) / len(eps_list), + markerfacecolor="none", + markeredgecolor=jacobian_colors[j_type], + ) + axs[i].set_title(f"FwdKine[{i}]") + axs[i].grid(True) + + fig.suptitle( + f"Fwd Kinematics components as a function of kappa\ns = {s:.3f}, sigma_x = {sigma_x:.3f}, sigma_y = {sigma_y - 1:.3f}" + ) + param_legend = 0.85 + fig.tight_layout(rect=[0, 0, param_legend, 1]) + + handles, labels = axs[0].get_legend_handles_labels() + # Supprimer doublons + unique = dict(zip(labels, handles)) + handles, labels = list(unique.values()), list(unique.keys()) + handles += [Patch(facecolor="white")] + labels += [f"s = {s:.2f}"] + if fig.legends: + for leg in fig.legends: + leg.remove() + fig.legend(handles, labels, loc="center left", bbox_to_anchor=(param_legend, 0.5)) + + +# === Valeurs initiales === +initial_s = float(params["l"][0] / 2) +initial_sigma_x = float(borne_sigma_x / 2) +initial_sigma_y = float(borne_sigma_y / 2) + +# === Première visualisation === +FwdKine_plot([eps_options[0]], initial_s, initial_sigma_x, initial_sigma_y, fig, axs) + +# === Sliders === +s_slider = Slider( + plt.axes([0.2, 0.01, 0.65, 0.03]), + "s", + float(s_values[0]), + float(s_values[-1]), + valinit=initial_s, +) +sigma_x_slider = Slider( + plt.axes([0.2, 0.05, 0.65, 0.03]), + "sigma_x", + float(sigma_x_values[0]), + float(sigma_x_values[-1]), + valinit=initial_sigma_x, +) +sigma_y_slider = Slider( + plt.axes([0.2, 0.09, 0.65, 0.03]), + "sigma_y", + float(sigma_y_values[0]), + float(sigma_y_values[-1]), + valinit=initial_sigma_y, +) +eps_slider = Slider( + plt.axes([0.2, 0.13, 0.65, 0.03]), + "eps (log scale)", + 0, + len(eps_options) - 1, + valinit=0, + valstep=1, +) + + +def on_slider_change(s_val, sigma_x_val, sigma_y_val): + s = float(s_val) + sigma_x = float(sigma_x_val) + sigma_y = float(sigma_y_val) + eps_val = get_eps_from_slider() + FwdKine_plot([eps_val], s, sigma_x, sigma_y, fig, axs) + fig.canvas.draw_idle() + + +def update_sliders(_): + on_slider_change(s_slider.val, sigma_x_slider.val, sigma_y_slider.val) + + +s_slider.on_changed(update_sliders) +sigma_x_slider.on_changed(update_sliders) +sigma_y_slider.on_changed(update_sliders) +eps_slider.on_changed(update_sliders) + +# === Boutons et CheckBoxes === +reset_ax = plt.axes([0.87, 0.6, 0.1, 0.05]) +reset_button = Button(reset_ax, "Reset sliders") +reset_button.on_clicked( + lambda event: ( + s_slider.reset(), + sigma_x_slider.reset(), + sigma_y_slider.reset(), + eps_slider.reset(), + ) +) + +check_ax = plt.axes([0.87, 0.7, 0.12, 0.15]) +check = CheckButtons(check_ax, jacobian_types, [True] * len(jacobian_types)) +check_ax.set_title("Jacobian types") + + +def on_check(label): + global list_of_type_of_jacobian + list_of_type_of_jacobian = [ + jacobian_types[i] for i, v in enumerate(check.get_status()) if v + ] + update_sliders(None) + + +check.on_clicked(on_check) + +plt.show() diff --git a/tests/gui_test_jacobian_eps_planar_pcs.py b/tests/gui_test_jacobian_eps_planar_pcs.py new file mode 100644 index 0000000..36aef83 --- /dev/null +++ b/tests/gui_test_jacobian_eps_planar_pcs.py @@ -0,0 +1,262 @@ +import jax.numpy as jnp +import matplotlib.pyplot as plt +from matplotlib.patches import Patch +from matplotlib.widgets import Slider, Button, CheckButtons +from matplotlib import rc + +rc("animation", html="html5") + +from jsrm.systems import planar_pcs_num, planar_pcs_sym +from pathlib import Path +import jsrm +from tqdm import tqdm +from matplotlib.animation import FFMpegWriter + +# === Paramètres initiaux === +num_segments = 1 +rho = 1070 * jnp.ones((num_segments,)) +params = { + "th0": jnp.array(0.0), + "l": 1e-1 * jnp.ones((num_segments,)), + "r": 2e-2 * jnp.ones((num_segments,)), + "rho": rho, + "g": jnp.array([0.0, 9.81]), + "E": 2e3 * jnp.ones((num_segments,)), + "G": 1e3 * jnp.ones((num_segments,)), +} +params["D"] = 1e-3 * jnp.diag( + ( + jnp.repeat(jnp.array([[1e0, 1e3, 1e3]]), num_segments, axis=0) + * params["l"][:, None] + ).flatten() +) +strain_selector = jnp.ones((3 * num_segments,), dtype=bool) + + +# === Chargement des fonctions Jacobiennes === +def get_jacobian_fn(jacobian_type): + if jacobian_type == "symbolic": + sym_exp_filepath = ( + Path(jsrm.__file__).parent + / "symbolic_expressions" + / f"planar_pcs_ns-{num_segments}.dill" + ) + _, _, _, aux = planar_pcs_sym.factory(sym_exp_filepath, strain_selector) + else: + _, _, _, aux = planar_pcs_num.factory( + num_segments, + strain_selector, + integration_type="gauss-legendre", + param_integration=5, + jacobian_type=jacobian_type, + ) + return aux["jacobian_fn"] + + +J_autodiff_fn = get_jacobian_fn("autodiff") +J_explicit_fn = get_jacobian_fn("explicit") +J_symbolic_fn = get_jacobian_fn("symbolic") + +jacobian_colors = {"symbolic": "green", "explicit": "orange", "autodiff": "blue"} +jacobian_markers = {"symbolic": "s", "explicit": "x", "autodiff": "o"} +jacobian_types = ["symbolic", "explicit", "autodiff"] +list_of_type_of_jacobian = jacobian_types.copy() + +# === eps discret === +eps_options = [None, 1e-6, 1e-5, 1e-4, 1e-3, 1e-2, 1e-1] +eps_labels = ["None"] + [f"1e-{i}" for i in range(6, 0, -1)] + + +def get_eps_from_slider(): + return eps_options[int(eps_slider.val)] + + +# === Variables physiques === +borne_kappa = 1e-1 +nb_kappa = 51 +borne_sigma_x, borne_sigma_y = 1e-1, 1e-1 +nb_sigma_x, nb_sigma_y, nb_s = 50, 50, 50 +sigma_x_values = jnp.linspace(0, borne_sigma_x, nb_sigma_x) +sigma_y_values = jnp.linspace(0, borne_sigma_y, nb_sigma_y) +s_values = jnp.linspace(0, float(params["l"][0]), nb_s) + +# === Tracé principal === +nb_colomns = 3 +nb_rows = 3 +fig, axs = plt.subplots(nb_colomns, nb_rows, figsize=(15, 8)) + + +def J_plot(eps_list, s, sigma_x, sigma_y, kappa_max, fig, axs): + kappa_values = jnp.linspace(-kappa_max, kappa_max, nb_kappa) + for ax_row in axs: + for ax in ax_row: + ax.clear() + ax.ticklabel_format(style="sci", axis="both", scilimits=(0, 0)) + ax.set_xlabel("kappa (bending strain)") + ax.set_ylabel("Jacobian components") + + for i_eps, eps in enumerate(eps_list): + J_auto, J_exp, J_symb = [], [], [] + for kappa in kappa_values: + q = jnp.array([kappa, sigma_x, sigma_y - 1.0] * num_segments) + J_auto.append(J_autodiff_fn(params, q, s, eps=eps)) + J_exp.append(J_explicit_fn(params, q, s, eps=eps)) + J_symb.append(J_symbolic_fn(params, q, s, eps=eps)) + J_auto, J_exp, J_symb = jnp.stack(J_auto), jnp.stack(J_exp), jnp.stack(J_symb) + + for i in range(nb_colomns): + for j in range(nb_rows): + if eps is not None: + axs[i, j].axvline( + eps, + color="red", + linestyle=":", + linewidth=2, + alpha=(i_eps + 1) / len(eps_list), + label=f"+/-eps={eps:.2e}", + ) + axs[i, j].axvline( + -eps, + color="red", + linestyle=":", + linewidth=2, + alpha=(i_eps + 1) / len(eps_list), + ) + + for j_type in list_of_type_of_jacobian: + J = {"symbolic": J_symb, "explicit": J_exp, "autodiff": J_auto}[ + j_type + ] + axs[i, j].plot( + kappa_values, + J[:, i, j], + marker=jacobian_markers[j_type], + label=f"J_{j_type}", + color=jacobian_colors[j_type], + alpha=(i_eps + 1) / len(eps_list), + markerfacecolor="none", + markeredgecolor=jacobian_colors[j_type], + ) + axs[i, j].set_title(f"J[{i}, {j}]") + axs[i, j].grid(True) + + fig.suptitle( + f"Jacobian components as a function of kappa\ns = {s:.3f}, sigma_x = {sigma_x:.3f}, sigma_y = {sigma_y - 1:.3f}" + ) + param_legend = 0.85 + fig.tight_layout(rect=[0, 0, param_legend, 1]) + + handles, labels = axs[0, 0].get_legend_handles_labels() + # Supprimer doublons + unique = dict(zip(labels, handles)) + handles, labels = list(unique.values()), list(unique.keys()) + handles += [Patch(facecolor="white")] + labels += [f"s = {s:.2f}"] + if fig.legends: + for leg in fig.legends: + leg.remove() + fig.legend(handles, labels, loc="center left", bbox_to_anchor=(param_legend, 0.5)) + + +# === Valeurs initiales === +initial_s = float(params["l"][0] / 2) +initial_sigma_x = float(borne_sigma_x / 2) +initial_sigma_y = float(borne_sigma_y / 2) +initial_borne_kappa = float(borne_kappa) + +# === Première visualisation === +J_plot( + [eps_options[0]], + initial_s, + initial_sigma_x, + initial_sigma_y, + initial_borne_kappa, + fig, + axs, +) + +# === Sliders === +s_slider = Slider( + plt.axes([0.2, 0.01, 0.65, 0.03]), + "s", + float(s_values[0]), + float(s_values[-1]), + valinit=initial_s, +) +sigma_x_slider = Slider( + plt.axes([0.2, 0.05, 0.65, 0.03]), + "sigma_x", + float(sigma_x_values[0]), + float(sigma_x_values[-1]), + valinit=initial_sigma_x, +) +sigma_y_slider = Slider( + plt.axes([0.2, 0.09, 0.65, 0.03]), + "sigma_y", + float(sigma_y_values[0]), + float(sigma_y_values[-1]), + valinit=initial_sigma_y, +) +eps_slider = Slider( + plt.axes([0.2, 0.13, 0.65, 0.03]), + "eps (log scale)", + 0, + len(eps_options) - 1, + valinit=0, + valstep=1, +) +kappa_slider = Slider( + plt.axes([0.2, 0.17, 0.65, 0.03]), "kappa max", 0.01, 0.5, valinit=borne_kappa +) + + +def on_slider_change(s_val, sigma_x_val, sigma_y_val, kappa_val): + s = float(s_val) + sigma_x = float(sigma_x_val) + sigma_y = float(sigma_y_val) + kappa_val = float(kappa_val) + eps_val = get_eps_from_slider() + J_plot([eps_val], s, sigma_x, sigma_y, kappa_val, fig, axs) + fig.canvas.draw_idle() + + +def update_sliders(_): + on_slider_change( + s_slider.val, sigma_x_slider.val, sigma_y_slider.val, kappa_slider.val + ) + + +s_slider.on_changed(update_sliders) +sigma_x_slider.on_changed(update_sliders) +sigma_y_slider.on_changed(update_sliders) +eps_slider.on_changed(update_sliders) +kappa_slider.on_changed(update_sliders) + +# === Boutons et CheckBoxes === +reset_ax = plt.axes([0.87, 0.6, 0.1, 0.05]) +reset_button = Button(reset_ax, "Reset sliders") +reset_button.on_clicked( + lambda event: ( + s_slider.reset(), + sigma_x_slider.reset(), + sigma_y_slider.reset(), + eps_slider.reset(), + ) +) + +check_ax = plt.axes([0.87, 0.7, 0.12, 0.15]) +check = CheckButtons(check_ax, jacobian_types, [True] * len(jacobian_types)) +check_ax.set_title("Jacobian types") + + +def on_check(label): + global list_of_type_of_jacobian + list_of_type_of_jacobian = [ + jacobian_types[i] for i, v in enumerate(check.get_status()) if v + ] + update_sliders(None) + + +check.on_clicked(on_check) + +plt.show() diff --git a/tests/test_planar_pcs.py b/tests/test_planar_pcs.py deleted file mode 100644 index 7998644..0000000 --- a/tests/test_planar_pcs.py +++ /dev/null @@ -1,127 +0,0 @@ -import jax - -jax.config.update("jax_enable_x64", True) # double precision -from jax import Array -from jax import numpy as jnp -import jsrm -from functools import partial -from numpy.testing import assert_allclose -from pathlib import Path - -from jsrm.systems import planar_pcs, euler_lagrangian -from jsrm.utils.tolerance import Tolerance - - -def constant_strain_inverse_kinematics_fn(params, xi_eq, chi, s) -> Array: - # split the chi vector into x, y, and th0 - px, py, th = chi - th0 = params["th0"].item() - print("th0 = ", th0) - # xi = (th - th0) / (2 * s) * jnp.array([ - # 2.0, - # (jnp.sin(th0)*px+jnp.cos(th0)*py) - (jnp.cos(th0)*px-jnp.sin(th0)*py)*jnp.sin(th-th0)/(jnp.cos(th-th0)-1), - # -(jnp.cos(th0)*px-jnp.sin(th0)*py) - (jnp.sin(th0)*px+jnp.cos(th0)*py)*jnp.sin(th-th0)/(jnp.cos(th-th0)-1) - # ]) - xi = (th - th0) / (2 * s) * jnp.array([ - 2.0, - (-jnp.sin(th0)*px+jnp.cos(th0)*py) - (jnp.cos(th0)*px+jnp.sin(th0)*py)*jnp.sin(th-th0)/(jnp.cos(th-th0)-1), - -(jnp.cos(th0)*px+jnp.sin(th0)*py) - (-jnp.sin(th0)*px+jnp.cos(th0)*py)*jnp.sin(th-th0)/(jnp.cos(th-th0)-1) - ]) - q = xi - xi_eq - return q - -def test_planar_cs(): - sym_exp_filepath = ( - Path(jsrm.__file__).parent / "symbolic_expressions" / "planar_pcs_ns-1.dill" - ) - params = { - "th0": jnp.array(0.0), # initial orientation angle [rad] - "l": jnp.array([1e-1]), - "r": jnp.array([2e-2]), - "rho": 1000 * jnp.ones((1,)), - "g": jnp.array([0.0, -9.81]), - "E": 1e8 * jnp.ones((1,)), # Elastic modulus [Pa] - "G": 1e7 * jnp.ones((1,)), # Shear modulus [Pa] - } - # activate all strains (i.e. bending, shear, and axial) - strain_selector = jnp.ones((3,), dtype=bool) - - xi_eq = jnp.array([0.0, 0.0, 1.0]) - strain_basis, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns = ( - planar_pcs.factory(sym_exp_filepath, strain_selector, xi_eq) - ) - forward_dynamics_fn = partial( - euler_lagrangian.forward_dynamics, dynamical_matrices_fn - ) - nonlinear_state_space_fn = partial( - euler_lagrangian.nonlinear_state_space, dynamical_matrices_fn - ) - - # test forward kinematics - assert_allclose( - forward_kinematics_fn(params, q=jnp.zeros((3,)), s=params["l"][0] / 2), - jnp.array([0.0, params["l"][0] / 2, 0.0]), - rtol=Tolerance.rtol(), - atol=Tolerance.atol(), - ) - assert_allclose( - forward_kinematics_fn(params, q=jnp.zeros((3,)), s=params["l"][0]), - jnp.array([0.0, params["l"][0], 0.0]), - rtol=Tolerance.rtol(), - atol=Tolerance.atol(), - ) - assert_allclose( - forward_kinematics_fn(params, q=jnp.array([0.0, 0.0, 1.0]), s=params["l"][0]), - jnp.array([0.0, 2 * params["l"][0], 0.0]), - rtol=Tolerance.rtol(), - atol=Tolerance.atol(), - ) - assert_allclose( - forward_kinematics_fn(params, q=jnp.array([0.0, 1.0, 0.0]), s=params["l"][0]), - params["l"][0] * jnp.array([1.0, 1.0, 0.0]), - rtol=Tolerance.rtol(), - atol=Tolerance.atol(), - ) - - # test inverse kinematics - params_ik = params.copy() - ik_th0_ls = [-jnp.pi / 2, -jnp.pi / 4, 0.0, jnp.pi / 4, jnp.pi / 2] - ik_q_ls = [ - jnp.array([0.1, 0.0, 0.0]), - jnp.array([0.1, 0.0, 0.2]), - jnp.array([0.1, 0.5, 0.1]), - jnp.array([1.0, 0.5, 0.2]), - jnp.array([-1.0, 0.0, 0.0]), - ] - for ik_th0 in ik_th0_ls: - params_ik["th0"] = jnp.array(ik_th0) - for q in ik_q_ls: - s = params["l"][0] - chi = forward_kinematics_fn(params_ik, q=q, s=s) - q_ik = constant_strain_inverse_kinematics_fn(params_ik, xi_eq, chi, s) - print("q = ", q, "q_ik = ", q_ik) - assert_allclose(q, q_ik, rtol=Tolerance.rtol(), atol=Tolerance.atol()) - - # test dynamical matrices - q, q_d = jnp.zeros((3,)), jnp.zeros((3,)) - B, C, G, K, D, A = dynamical_matrices_fn(params, q, q_d) - assert_allclose(K, jnp.zeros((3,))) - assert_allclose( - A, - jnp.eye(3), - ) - - q = jnp.array([jnp.pi / (2 * params["l"][0]), 0.0, 0.0]) - q_d = jnp.zeros((3,)) - B, C, G, K, D, alpha = dynamical_matrices_fn(params, q, q_d) - - print("B =\n", B) - print("C =\n", C) - print("G =\n", G) - print("K =\n", K) - print("D =\n", D) - print("alpha =\n", alpha) - - -if __name__ == "__main__": - test_planar_cs() diff --git a/tests/test_planar_pcs_num.py b/tests/test_planar_pcs_num.py new file mode 100644 index 0000000..07323d2 --- /dev/null +++ b/tests/test_planar_pcs_num.py @@ -0,0 +1,283 @@ +import jax + +jax.config.update("jax_enable_x64", True) # double precision +from jax import Array +from jax import numpy as jnp +import jsrm +from functools import partial +from numpy.testing import assert_allclose +from pathlib import Path + +from jsrm.systems import planar_pcs_num, euler_lagrangian +from jsrm.utils.tolerance import Tolerance + +from typing import Optional, Literal + + +def constant_strain_inverse_kinematics_fn(params, xi_eq, chi, s) -> Array: + # split the chi vector into x, y, and th0 + px, py, th = chi + th0 = params["th0"].item() + print("th0 = ", th0) + xi = ( + (th - th0) + / (2 * s) + * jnp.array( + [ + 2.0, + (-jnp.sin(th0) * px + jnp.cos(th0) * py) + - (jnp.cos(th0) * px + jnp.sin(th0) * py) + * jnp.sin(th - th0) + / (jnp.cos(th - th0) - 1), + -(jnp.cos(th0) * px + jnp.sin(th0) * py) + - (-jnp.sin(th0) * px + jnp.cos(th0) * py) + * jnp.sin(th - th0) + / (jnp.cos(th - th0) - 1), + ] + ) + ) + q = xi - xi_eq + return q + + +def test_planar_cs_num( + type_of_integration: Optional[ + Literal["gauss-legendre", "gauss-kronrad", "trapezoid"] + ] = "gauss-legendre", + type_of_jacobian: Optional[Literal["explicit", "autodiff"]] = "explicit", +): + """ + Test the planar constant strain system with numerical integration and Jacobian for 1 segment. + + Args: + type_of_integration (Literal["gauss-legendre", "gauss-kronrad", "trapezoid"], optional): + Type of integration method to use. + "gauss-kronrad" for Gauss-Kronrad rule, "gauss-legendre" for Gauss-Legendre rule, + "trapezoid" for trapezoid rule. + Defaults to "gauss-legendre". + type_of_jacobian (Literal["explicit", "autodiff"], optional): + Type of Jacobian method to use. + "explicit" for explicit Jacobian, "autodiff" for automatic differentiation. + Defaults to "explicit". + """ + params = { + "th0": jnp.array(0.0), # initial orientation angle [rad] + "l": jnp.array([1e-1]), + "r": jnp.array([2e-2]), + "rho": 1000 * jnp.ones((1,)), + "g": jnp.array([0.0, -9.81]), + "E": 1e8 * jnp.ones((1,)), # Elastic modulus [Pa] + "G": 1e7 * jnp.ones((1,)), # Shear modulus [Pa] + } + # activate all strains (i.e. bending, shear, and axial) + strain_selector = jnp.ones((3,), dtype=bool) + + xi_eq = jnp.array([0.0, 0.0, 1.0]) + + num_segments = 1 + if type_of_integration == "gauss-kronrad": + # use Gauss-Kronrad rule for integration + param_integration = 15 + elif type_of_integration == "gauss-legendre": + # use Gauss-Legendre rule for integration + param_integration = 5 + elif type_of_integration == "trapezoid": + # use trapezoid rule for integration + param_integration = 1000 + strain_basis, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns = ( + planar_pcs_num.factory( + num_segments, + strain_selector, + integration_type=type_of_integration, + param_integration=param_integration, + jacobian_type=type_of_jacobian, + ) + ) + forward_dynamics_fn = partial( + euler_lagrangian.forward_dynamics, dynamical_matrices_fn + ) + nonlinear_state_space_fn = partial( + euler_lagrangian.nonlinear_state_space, dynamical_matrices_fn + ) + + # ======================================== + # Test of the functions + # ======================================== + + # test forward kinematics + print("\nTesting forward kinematics... ------------------------") + test_cases = [ + ( + jnp.zeros((3,)), + params["l"][0] / 2, + jnp.array([0.0, params["l"][0] / 2, 0.0]), + ), + (jnp.zeros((3,)), params["l"][0], jnp.array([0.0, params["l"][0], 0.0])), + ( + jnp.array([0.0, 0.0, 1.0]), + params["l"][0], + jnp.array([0.0, 2 * params["l"][0], 0.0]), + ), + ( + jnp.array([0.0, 1.0, 0.0]), + params["l"][0], + params["l"][0] * jnp.array([1.0, 1.0, 0.0]), + ), + ] + + for q, s, expected in test_cases: + print("q = ", q, "s = ", s) + chi = forward_kinematics_fn(params, q=q, s=s) + assert not jnp.isnan(chi).any(), "Forward kinematics output contains NaN!" + assert_allclose(chi, expected, rtol=Tolerance.rtol(), atol=Tolerance.atol()) + print("[Valid test]\n") + + # test inverse kinematics + print("\nTesting inverse kinematics... ------------------------") + params_ik = params.copy() + ik_th0_ls = [-jnp.pi / 2, -jnp.pi / 4, 0.0, jnp.pi / 4, jnp.pi / 2] + ik_q_ls = [ + jnp.array([0.1, 0.0, 0.0]), + jnp.array([0.1, 0.0, 0.2]), + jnp.array([0.1, 0.5, 0.1]), + jnp.array([1.0, 0.5, 0.2]), + jnp.array([-1.0, 0.0, 0.0]), + ] + for ik_th0 in ik_th0_ls: + params_ik["th0"] = jnp.array(ik_th0) + for q in ik_q_ls: + s = params_ik["l"][0] + print("q = ", q, "s = ", s, "th0 = ", ik_th0) + chi = forward_kinematics_fn(params_ik, q=q, s=s) + assert not jnp.isnan(chi).any(), "Forward kinematics output contains NaN!" + q_ik = constant_strain_inverse_kinematics_fn(params_ik, xi_eq, chi, s) + assert not jnp.isnan(q_ik).any(), "Inverse kinematics output contains NaN!" + assert_allclose(q, q_ik, rtol=Tolerance.rtol(), atol=Tolerance.atol()) + print("[Valid test]\n") + + # test dynamical matrices + print("\nTesting dynamical matrices... ------------------------") + q = jnp.zeros((3,)) + q_d = jnp.zeros((3,)) + print("q = ", q, "q_d = ", q_d) + B, C, G, K, D, A = dynamical_matrices_fn(params, q, q_d) + assert not jnp.isnan(B).any(), "B matrix contains NaN!" + assert not jnp.isnan(C).any(), "C matrix contains NaN!" + assert not jnp.isnan(G).any(), "G matrix contains NaN!" + assert not jnp.isnan(K).any(), "K matrix contains NaN!" + assert not jnp.isnan(D).any(), "D matrix contains NaN!" + assert not jnp.isnan(A).any(), "A matrix contains NaN!" + print("testing K") + assert_allclose(K, jnp.zeros((3,))) + print("[Valid test]\n") + print("testing A") + assert_allclose( + A, + jnp.eye(3), + ) + print("[Valid test]\n") + + q = jnp.array([jnp.pi / (2 * params["l"][0]), 0.0, 0.0]) + q_d = jnp.zeros((3,)) + print("q = ", q, "q_d = ", q_d) + B, C, G, K, D, A = dynamical_matrices_fn(params, q, q_d) + assert not jnp.isnan(B).any(), "B matrix contains NaN!" + assert not jnp.isnan(C).any(), "C matrix contains NaN!" + assert not jnp.isnan(G).any(), "G matrix contains NaN!" + assert not jnp.isnan(K).any(), "K matrix contains NaN!" + assert not jnp.isnan(D).any(), "D matrix contains NaN!" + assert not jnp.isnan(A).any(), "A matrix contains NaN!" + + print("B =\n", B) + print("C =\n", C) + print("G =\n", G) + print("K =\n", K) + print("D =\n", D) + print("A =\n", A) + print("[To check]") + + # test energies + print("\nTesting energies... ------------------------") + kinetic_energy_fn = auxiliary_fns["kinetic_energy_fn"] + potential_energy_fn = auxiliary_fns["potential_energy_fn"] + + q = jnp.zeros((3,)) + q_d = jnp.zeros((3,)) + print("q = ", q, "q_d = ", q_d) + + print("Testing kinetic energy...") + E_kin = kinetic_energy_fn(params, q, q_d) + assert not jnp.isnan(E_kin).any(), "Kinetic energy contains NaN!" + E_kin_th = 0.0 + assert_allclose(E_kin, E_kin_th, rtol=Tolerance.rtol(), atol=Tolerance.atol()) + print("[Valid test]\n") + + print("Testing potential energy...") + E_pot = potential_energy_fn(params, q) + assert not jnp.isnan(E_pot).any(), "Potential energy contains NaN!" + E_pot_th = jnp.array( + 0.5 + * params["rho"][0] + * jnp.pi + * params["r"][0] ** 2 + * jnp.linalg.norm(params["g"]) + * params["l"][0] ** 2 + ) + assert_allclose(E_pot, E_pot_th, rtol=Tolerance.rtol(), atol=Tolerance.atol()) + print("[Valid test]\n") + + # test jacobian + print("\nTesting jacobian... ------------------------") + jacobian_fn = auxiliary_fns["jacobian_fn"] + chi = forward_kinematics_fn(params, q=q, s=params["l"][0]) + print("q = ", q, "s = ", params["l"][0]) + J = jacobian_fn(params, q, s=params["l"][0]) + assert not jnp.isnan(J).any(), "Jacobian contains NaN!" + print("Jacobian J =\n", J) + # Test the differential relation: delta_chi ≈ J * delta_q + print("Testing differential relation: delta_chi ≈ J * delta_q") + delta_q = jnp.array([1e-6, -1e-6, 2e-6]) + chi_plus = forward_kinematics_fn(params, q=q + delta_q, s=params["l"][0]) + chi_pred = chi + J @ delta_q + assert_allclose(chi_plus, chi_pred, rtol=Tolerance.rtol(), atol=Tolerance.atol()) + print("[Valid test]\n") + + # test forward dynamics + print("\nTesting forward dynamics... ------------------------") + q = jnp.zeros((3,)) + q_d = jnp.zeros((3,)) + tau = jnp.array([0.0, 0.0, 0.0]) # no external forces + params_bis = params.copy() + params_bis["g"] = jnp.array([0.0, 0.0]) # no gravity for this test + print("q = ", q, "q_d = ", q_d, "tau = ", tau, "g = ", params_bis["g"]) + q_dd = forward_dynamics_fn(params_bis, q, q_d, tau) + assert not jnp.isnan(q_dd).any(), "Forward dynamics output contains NaN!" + assert_allclose(q_dd, jnp.zeros((3,)), rtol=Tolerance.rtol(), atol=Tolerance.atol()) + print("[Valid test]\n") + + # test nonlinear state space + print("\nTesting nonlinear state space... ------------------------") + x = jnp.concatenate([q, q_d]) + print("x = ", x, "tau = ", tau) + x_dot = nonlinear_state_space_fn(params, x, tau) + assert not jnp.isnan(x_dot).any(), "Nonlinear state space output contains NaN!" + print("x_dot = ", x_dot) + print("[To check]") + + +if __name__ == "__main__": + list_of_integration_types = ["gauss-legendre", "gauss-kronrad", "trapezoid"] + list_of_jacobian_types = ["autodiff", "explicit"] + + for integration_type in list_of_integration_types: + for jacobian_type in list_of_jacobian_types: + print( + "\n================================================================================================" + ) + print( + f"Testing {integration_type} integration with {jacobian_type} Jacobian..." + ) + test_planar_cs_num( + type_of_integration=integration_type, + type_of_jacobian=jacobian_type, + ) diff --git a/tests/test_planar_pcs_sym.py b/tests/test_planar_pcs_sym.py new file mode 100644 index 0000000..ec4409d --- /dev/null +++ b/tests/test_planar_pcs_sym.py @@ -0,0 +1,234 @@ +import jax + +jax.config.update("jax_enable_x64", True) # double precision +from jax import Array +from jax import numpy as jnp +import jsrm +from functools import partial +from numpy.testing import assert_allclose +from pathlib import Path + +from jsrm.systems import euler_lagrangian, planar_pcs_sym +from jsrm.utils.tolerance import Tolerance + + +def constant_strain_inverse_kinematics_fn(params, xi_eq, chi, s) -> Array: + # split the chi vector into x, y, and th0 + px, py, th = chi + th0 = params["th0"].item() + print("th0 = ", th0) + xi = ( + (th - th0) + / (2 * s) + * jnp.array( + [ + 2.0, + (-jnp.sin(th0) * px + jnp.cos(th0) * py) + - (jnp.cos(th0) * px + jnp.sin(th0) * py) + * jnp.sin(th - th0) + / (jnp.cos(th - th0) - 1), + -(jnp.cos(th0) * px + jnp.sin(th0) * py) + - (-jnp.sin(th0) * px + jnp.cos(th0) * py) + * jnp.sin(th - th0) + / (jnp.cos(th - th0) - 1), + ] + ) + ) + q = xi - xi_eq + return q + + +def test_planar_cs(): + sym_exp_filepath = ( + Path(jsrm.__file__).parent / "symbolic_expressions" / "planar_pcs_ns-1.dill" + ) + params = { + "th0": jnp.array(0.0), # initial orientation angle [rad] + "l": jnp.array([1e-1]), + "r": jnp.array([2e-2]), + "rho": 1000 * jnp.ones((1,)), + "g": jnp.array([0.0, -9.81]), + "E": 1e8 * jnp.ones((1,)), # Elastic modulus [Pa] + "G": 1e7 * jnp.ones((1,)), # Shear modulus [Pa] + } + # activate all strains (i.e. bending, shear, and axial) + strain_selector = jnp.ones((3,), dtype=bool) + + xi_eq = jnp.array([0.0, 0.0, 1.0]) + strain_basis, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns = ( + planar_pcs_sym.factory(sym_exp_filepath, strain_selector, xi_eq) + ) + forward_dynamics_fn = partial( + euler_lagrangian.forward_dynamics, dynamical_matrices_fn + ) + nonlinear_state_space_fn = partial( + euler_lagrangian.nonlinear_state_space, dynamical_matrices_fn + ) + + # ======================================== + # Test of the functions + # ======================================== + + # test forward kinematics + print("\nTesting forward kinematics... ------------------------") + test_cases = [ + ( + jnp.zeros((3,)), + params["l"][0] / 2, + jnp.array([0.0, params["l"][0] / 2, 0.0]), + ), + (jnp.zeros((3,)), params["l"][0], jnp.array([0.0, params["l"][0], 0.0])), + ( + jnp.array([0.0, 0.0, 1.0]), + params["l"][0], + jnp.array([0.0, 2 * params["l"][0], 0.0]), + ), + ( + jnp.array([0.0, 1.0, 0.0]), + params["l"][0], + params["l"][0] * jnp.array([1.0, 1.0, 0.0]), + ), + ] + + for q, s, expected in test_cases: + print("q = ", q, "s = ", s) + chi = forward_kinematics_fn(params, q=q, s=s) + assert not jnp.isnan(chi).any(), "Forward kinematics output contains NaN!" + assert_allclose(chi, expected, rtol=Tolerance.rtol(), atol=Tolerance.atol()) + print("[Valid test]\n") + + # test inverse kinematics + print("\nTesting inverse kinematics... ------------------------") + params_ik = params.copy() + ik_th0_ls = [-jnp.pi / 2, -jnp.pi / 4, 0.0, jnp.pi / 4, jnp.pi / 2] + ik_q_ls = [ + jnp.array([0.1, 0.0, 0.0]), + jnp.array([0.1, 0.0, 0.2]), + jnp.array([0.1, 0.5, 0.1]), + jnp.array([1.0, 0.5, 0.2]), + jnp.array([-1.0, 0.0, 0.0]), + ] + for ik_th0 in ik_th0_ls: + params_ik["th0"] = jnp.array(ik_th0) + for q in ik_q_ls: + s = params_ik["l"][0] + print("q = ", q, "s = ", s, "th0 = ", ik_th0) + chi = forward_kinematics_fn(params_ik, q=q, s=s) + assert not jnp.isnan(chi).any(), "Forward kinematics output contains NaN!" + q_ik = constant_strain_inverse_kinematics_fn(params_ik, xi_eq, chi, s) + assert not jnp.isnan(q_ik).any(), "Inverse kinematics output contains NaN!" + assert_allclose(q, q_ik, rtol=Tolerance.rtol(), atol=Tolerance.atol()) + print("[Valid test]\n") + + # test dynamical matrices + print("\nTesting dynamical matrices... ------------------------") + q = jnp.zeros((3,)) + q_d = jnp.zeros((3,)) + print("q = ", q, "q_d = ", q_d) + B, C, G, K, D, A = dynamical_matrices_fn(params, q, q_d) + assert not jnp.isnan(B).any(), "B matrix contains NaN!" + assert not jnp.isnan(C).any(), "C matrix contains NaN!" + assert not jnp.isnan(G).any(), "G matrix contains NaN!" + assert not jnp.isnan(K).any(), "K matrix contains NaN!" + assert not jnp.isnan(D).any(), "D matrix contains NaN!" + assert not jnp.isnan(A).any(), "A matrix contains NaN!" + print("testing K") + assert_allclose(K, jnp.zeros((3,))) + print("[Valid test]\n") + print("testing A") + assert_allclose( + A, + jnp.eye(3), + ) + print("[Valid test]\n") + + q = jnp.array([jnp.pi / (2 * params["l"][0]), 0.0, 0.0]) + q_d = jnp.zeros((3,)) + print("q = ", q, "q_d = ", q_d) + B, C, G, K, D, A = dynamical_matrices_fn(params, q, q_d) + assert not jnp.isnan(B).any(), "B matrix contains NaN!" + assert not jnp.isnan(C).any(), "C matrix contains NaN!" + assert not jnp.isnan(G).any(), "G matrix contains NaN!" + assert not jnp.isnan(K).any(), "K matrix contains NaN!" + assert not jnp.isnan(D).any(), "D matrix contains NaN!" + assert not jnp.isnan(A).any(), "A matrix contains NaN!" + + print("B =\n", B) + print("C =\n", C) + print("G =\n", G) + print("K =\n", K) + print("D =\n", D) + print("A =\n", A) + print("[To check]") + + # test energies + print("\nTesting energies... ------------------------") + kinetic_energy_fn = auxiliary_fns["kinetic_energy_fn"] + potential_energy_fn = auxiliary_fns["potential_energy_fn"] + + q = jnp.zeros((3,)) + q_d = jnp.zeros((3,)) + print("q = ", q, "q_d = ", q_d) + + print("Testing kinetic energy...") + E_kin = kinetic_energy_fn(params, q, q_d) + assert not jnp.isnan(E_kin).any(), "Kinetic energy contains NaN!" + E_kin_th = 0.0 + assert_allclose(E_kin, E_kin_th, rtol=Tolerance.rtol(), atol=Tolerance.atol()) + print("[Valid test]\n") + + print("Testing potential energy...") + E_pot = potential_energy_fn(params, q) + assert not jnp.isnan(E_pot).any(), "Potential energy contains NaN!" + E_pot_th = jnp.array( + 0.5 + * params["rho"][0] + * jnp.pi + * params["r"][0] ** 2 + * jnp.linalg.norm(params["g"]) + * params["l"][0] ** 2 + ) + assert_allclose(E_pot, E_pot_th, rtol=Tolerance.rtol(), atol=Tolerance.atol()) + print("[Valid test]\n") + + # test jacobian + print("\nTesting jacobian... ------------------------") + jacobian_fn = auxiliary_fns["jacobian_fn"] + chi = forward_kinematics_fn(params, q=q, s=params["l"][0]) + print("q = ", q, "s = ", params["l"][0]) + J = jacobian_fn(params, q, s=params["l"][0]) + assert not jnp.isnan(J).any(), "Jacobian contains NaN!" + print("Jacobian J =\n", J) + # Test the differential relation: delta_chi ≈ J * delta_q + print("Testing differential relation: delta_chi ≈ J * delta_q") + delta_q = jnp.array([1e-6, -1e-6, 2e-6]) + chi_plus = forward_kinematics_fn(params, q=q + delta_q, s=params["l"][0]) + chi_pred = chi + J @ delta_q + assert_allclose(chi_plus, chi_pred, rtol=Tolerance.rtol(), atol=Tolerance.atol()) + print("[Valid test]\n") + + # test forward dynamics + print("\nTesting forward dynamics... ------------------------") + q = jnp.zeros((3,)) + q_d = jnp.zeros((3,)) + tau = jnp.array([0.0, 0.0, 0.0]) # no external forces + params_bis = params.copy() + params_bis["g"] = jnp.array([0.0, 0.0]) # no gravity for this test + print("q = ", q, "q_d = ", q_d, "tau = ", tau, "g = ", params_bis["g"]) + q_dd = forward_dynamics_fn(params_bis, q, q_d, tau) + assert not jnp.isnan(q_dd).any(), "Forward dynamics output contains NaN!" + assert_allclose(q_dd, jnp.zeros((3,)), rtol=Tolerance.rtol(), atol=Tolerance.atol()) + print("[Valid test]\n") + + # test nonlinear state space + print("\nTesting nonlinear state space... ------------------------") + x = jnp.concatenate([q, q_d]) + print("x = ", x, "tau = ", tau) + x_dot = nonlinear_state_space_fn(params, x, tau) + assert not jnp.isnan(x_dot).any(), "Nonlinear state space output contains NaN!" + print("x_dot = ", x_dot) + print("[To check]") + + +if __name__ == "__main__": + test_planar_cs()