Fix bug in implementation of potential_energy_fn for planar_pcs system

Author: mstoelzle

pyproject.toml

@@ -17,7 +17,7 @@ name = "jsrm"  # Required
 #
 # For a discussion on single-sourcing the version, see
 # https://packaging.python.org/guides/single-sourcing-package-version/
-version = "0.0.8"  # Required
+version = "0.0.9"  # Required
 
 # This is a one-line description or tagline of what your project does. This
 # corresponds to the "Summary" metadata field:

src/jsrm/systems/planar_pcs.py

@@ -111,6 +111,9 @@ def select_params_for_lambdify(params: Dict[str, Array]) -> List[Array]:
     G_lambda = sp.lambdify(
         params_syms_cat + sym_exps["state_syms"]["xi"], sym_exps["exps"]["G"], "jax"
     )
+    U_lambda = sp.lambdify(
+        params_syms_cat + sym_exps["state_syms"]["xi"], sym_exps["exps"]["U"], "jax"
+    )
 
     compute_stiffness_matrix_for_all_segments_fn = vmap(
         compute_planar_stiffness_matrix, in_axes=(0, 0, 0, 0), out_axes=0
@@ -300,9 +303,8 @@ def potential_energy_fn(params: Dict[str, Array], q: Array, eps: float = 1e4 * g
         U_K = (xi - xi_eq).T @ K @ (xi - xi_eq)  # evaluate K(xi) = K @ xi
 
         # gravitational potential energy
-        U_G = sp.Matrix([[0]])
         params_for_lambdify = select_params_for_lambdify(params)
-        U_G = G_lambda(*params_for_lambdify, *xi_epsed).squeeze() @ xi_epsed
+        U_G = U_lambda(*params_for_lambdify, *xi_epsed)
 
         # total potential energy
         U = (U_G + U_K).squeeze()