Fix jit decorators to apply JIT-compilation to

last level

Author: solangegbv

examples/benchmark_planar_pcs_num.py examples/benchmark_planar_pcs.pyexamples/benchmark_planar_pcs_num.py renamed to examples/benchmark_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)

examples/simulate_planar_pcs.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,

src/jsrm/integration.py

@@ -1,5 +1,5 @@
 from jax import numpy as jnp
-from jax import Array, lax, jit
+from jax import Array, lax
 
 def blk_diag(
     a: Array

src/jsrm/math_utils.py

@@ -13,7 +13,6 @@
 )
 from jsrm.math_utils import blk_diag
 
-
 def factory(
     filepath: Union[str, Path],
     strain_selector: Array = None,
@@ -54,7 +53,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
@@ -147,7 +145,6 @@ def select_params_for_lambdify_fn(params: Dict[str, Array]) -> List[Array]:
         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
@@ -263,7 +260,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:
@@ -297,7 +293,6 @@ def forward_kinematics_fn(
 
         return chi
 
-    @jit
     def jacobian_fn(
         params: Dict[str, Array], q: Array, s: Array, eps: float = global_eps
     ) -> Array:
@@ -333,7 +328,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]: