Merge branch 'main' into gvs

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(
@@ -158,6 +160,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,6 +218,7 @@ 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)

src/jsrm/math_utils.py

@@ -69,37 +69,4 @@ def blk_concat(
 
     c = blk_concat(a)
     print("Concatenated matrix:")
-    print(c)
-    
-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
+    print(c)