Correction for merging

Author: solangegbv

examples/simulate_planar_pcs_sym.py

@@ -150,14 +150,7 @@ def draw_robot(
     # the evolution of the generalized coordinates
     q_ts = sol.ys[:, :n_q]
     # 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)
-    
+    q_d_ts = sol.ys[:, n_q:]    
 
     s_max = jnp.array([jnp.sum(params["l"])])
 
@@ -167,7 +160,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):
@@ -225,11 +217,9 @@ 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,11 +1,6 @@
 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:
@@ -78,7 +73,8 @@ 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,

src/jsrm/systems/utils.py

@@ -9,14 +9,9 @@
 from jax import Array
 from typing import Dict, List, 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, Array]
-    sym_exps: Dict, params: Dict[str, Array]
 ) -> Dict:
     """
     Substitute robot parameters into symbolic expressions.
@@ -57,7 +52,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 +92,6 @@ def concatenate_params_syms(
 
 def compute_strain_basis(
     strain_selector: Array,
-) -> Array:
 ) -> Array:
     """
     Compute strain basis based on boolean strain selector.
@@ -124,9 +117,6 @@ def compute_strain_basis(
     return strain_basis
 
 
-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: