Merge branch 'main' into numerical-derivation

Author: solangegbv

examples/simulate_planar_pcs_sym.py

@@ -131,6 +131,8 @@ def draw_robot(
     ode_fn = ode_factory(dynamical_matrices_fn, params, tau)
     # jit the ODE function
     ode_fn = jax.jit(ode_fn)
+    # jit the ODE function
+    ode_fn = jax.jit(ode_fn)
     term = ODETerm(ode_fn)
 
     sol = diffeqsolve(
@@ -150,6 +152,13 @@ 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)
+    
+
     s_max = jnp.array([jnp.sum(params["l"])])
 
     forward_kinematics_fn_end_effector = partial(forward_kinematics_fn, params, s=s_max)
@@ -158,6 +167,7 @@ def draw_robot(
 
     # evaluate the forward kinematics along the trajectory
     chi_ee_ts = forward_kinematics_fn_end_effector(q_ts)
+    chi_ee_ts = forward_kinematics_fn_end_effector(q_ts)
     # plot the configuration vs time
     plt.figure()
     for segment_idx in range(num_segments):
@@ -215,9 +225,11 @@ def draw_robot(
     # plot the energy along the trajectory
     kinetic_energy_fn_vmapped = vmap(
         partial(jax.jit(auxiliary_fns["kinetic_energy_fn"]), params)
+        partial(jax.jit(auxiliary_fns["kinetic_energy_fn"]), params)
     )
     potential_energy_fn_vmapped = vmap(
         partial(jax.jit(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)

src/jsrm/math_utils.py

@@ -1,12 +1,19 @@
 from jax import numpy as jnp
 from jax import Array, lax
 
+def blk_diag(
+    a: Array
+) -> Array:
+from jax import Array, lax
+
 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)
 
@@ -33,6 +40,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]), dtype=a.dtype)
+    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],