GVS compatible with main branch

Author: solangegbv

examples/simulate_planar_pcs_sym.py

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

@@ -57,6 +57,40 @@ def blk_concat(
     b = a.transpose(1, 0, 2).reshape(a.shape[1], -1)
     return b
 
+    
+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
+
 if __name__ == "__main__":
     # Example usage
     a = jnp.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])

src/jsrm/systems/utils.py

@@ -16,7 +16,6 @@
 
 def substitute_params_into_all_symbolic_expressions(
     sym_exps: Dict, params: Dict[str, Array]
-    sym_exps: Dict, params: Dict[str, Array]
 ) -> Dict:
     """
     Substitute robot parameters into symbolic expressions.
@@ -57,7 +56,6 @@ def substitute_params_into_single_symbolic_expression(
     sym_exp: sp.Expr,
     params_syms: Dict[str, List[sp.Symbol]],
     params: Dict[str, Array],
-    params: Dict[str, Array],
 ) -> sp.Expr:
     """
     Substitute robot parameters into a single symbolic expression.
@@ -98,7 +96,6 @@ def concatenate_params_syms(
 
 def compute_strain_basis(
     strain_selector: Array,
-) -> Array:
 ) -> Array:
     """
     Compute strain basis based on boolean strain selector.