Add kinetic and potential energy functions for planar_pcs system

examples/simulate_planar_pcs.py

@@ -1,10 +1,12 @@
 import cv2  # importing cv2
+from functools import partial
 import jax
 
 jax.config.update("jax_enable_x64", True)  # double precision
 from diffrax import diffeqsolve, Dopri5, Euler, ODETerm, SaveAt
 from jax import Array, vmap
 from jax import numpy as jnp
+import matplotlib.pyplot as plt
 import numpy as onp
 from pathlib import Path
 from typing import Callable, Dict
@@ -42,6 +44,8 @@
 
 # define initial configuration
 q0 = jnp.array([10 * jnp.pi])
+# number of generalized coordinates
+n_q = q0.shape[0]
 
 # set simulation parameters
 dt = 1e-3  # time step
@@ -91,7 +95,7 @@ def draw_robot(
 
 
 if __name__ == "__main__":
-    strain_basis, forward_kinematics_fn, dynamical_matrices_fn = planar_pcs.factory(
+    strain_basis, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns = planar_pcs.factory(
         sym_exp_filepath, strain_selector
     )
     batched_forward_kinematics = vmap(
@@ -133,6 +137,25 @@ def draw_robot(
     )
 
     print("sol.ys =\n", sol.ys)
+    # 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:]
+
+    # plot the energy along the trajectory
+    kinetic_energy_fn_vmapped = vmap(partial(auxiliary_fns["kinetic_energy_fn"], params))
+    potential_energy_fn_vmapped = vmap(partial(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)
+    plt.figure()
+    plt.plot(video_ts, U_ts, label="Potential energy")
+    plt.plot(video_ts, T_ts, label="Kinetic energy")
+    plt.xlabel("Time [s]")
+    plt.ylabel("Energy [J]")
+    plt.legend()
+    plt.grid(True)
+    plt.box(True)
+    plt.show()
 
     # create video
     fourcc = cv2.VideoWriter_fourcc(*"MP4V")

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.7"  # Required
+version = "0.0.8"  # Required
 
 # This is a one-line description or tagline of what your project does. This
 # corresponds to the "Summary" metadata field:

src/jsrm/symbolic_derivation/planar_pcs.py

@@ -169,9 +169,10 @@ def symbolically_derive_planar_pcs_model(
             ),  # expression for end-effector pose of shape (3, )
             "J_sms": J_sms,
             "Jee": J_sms[-1].subs(s, l[-1]),
-            "B": B,
-            "C": C,
-            "G": G,
+            "B": B,  # mass matrix
+            "C": C,  # coriolis matrix
+            "G": G,  # gravity vector
+            "U": U,  # gravitational potential energy
         },
     }
 

src/jsrm/systems/planar_pcs.py

@@ -25,6 +25,7 @@ def factory(
         [Dict[str, Array], Array, Array],
         Tuple[Array, Array, Array, Array, Array, Array],
     ],
+    Dict[str, Callable],
 ]:
     """
     Create jax functions from file containing symbolic expressions.
@@ -38,6 +39,7 @@ def factory(
         B_xi: strain basis matrix of shape (3 * num_segments, n_q)
         forward_kinematics_fn: function that returns the p vector of shape (3, n_q) with the positions
         dynamical_matrices_fn: function that returns the B, C, G, K, D, and alpha matrices
+        auxiliary_fns: dictionary with auxiliary functions
     """
     # load saved symbolic data
     sym_exps = dill.load(open(str(filepath), "rb"))
@@ -249,5 +251,85 @@ def dynamical_matrices_fn(
         alpha = B_xi.T @ jnp.identity(n_xi) @ B_xi
 
         return B, C, G, K, D, alpha
+
 
-    return B_xi, forward_kinematics_fn, dynamical_matrices_fn
+    def kinetic_energy_fn(params: Dict[str, Array], q: Array, q_d: Array) -> Array:
+        """
+        Compute the kinetic energy of the system.
+        Args:
+            params: Dictionary of robot parameters
+            q: generalized coordinates of shape (n_q, )
+            q_d: generalized velocities of shape (n_q, )
+        Returns:
+            T: kinetic energy of shape ()
+        """
+        B, C, G, K, D, alpha = dynamical_matrices_fn(params, q=q, q_d=q_d)
+
+        # kinetic energy
+        T = (0.5 * q_d.T @ B @ q_d).squeeze()
+
+        return T
+
+
+    def potential_energy_fn(params: Dict[str, Array], q: Array, eps: float = 1e4 * global_eps) -> Array:
+        """
+        Compute the potential energy of the system.
+        Args:
+            params: Dictionary of robot parameters
+            q: generalized coordinates of shape (n_q, )
+            eps: small number to avoid singularities (e.g., division by zero)
+        Returns:
+            U: potential energy of shape ()
+        """
+        # map the configuration to the strains
+        xi = xi_eq + B_xi @ q
+        # add a small number to the bending strain to avoid singularities
+        xi_epsed = apply_eps_to_bend_strains(xi, eps)
+
+        # cross-sectional area and second moment of area
+        A = jnp.pi * params["r"] ** 2
+        Ib = A**2 / (4 * jnp.pi)
+
+        # elastic and shear modulus
+        E, G = params["E"], params["G"]
+        # stiffness matrix of shape (num_segments, 3, 3)
+        S = compute_stiffness_matrix_for_all_segments_fn(A, Ib, E, G)
+        # we define the elastic matrix of shape (n_xi, n_xi) as K(xi) = K @ xi where K is equal to
+        K = blk_diag(S)
+        # elastic energy
+        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
+
+        # total potential energy
+        U = (U_G + U_K).squeeze()
+
+        return U
+
+    def energy_fn(params: Dict[str, Array], q: Array, q_d: Array) -> Array:
+        """
+        Compute the total energy of the system.
+        Args:
+            params: Dictionary of robot parameters
+            q: generalized coordinates of shape (n_q, )
+            q_d: generalized velocities of shape (n_q, )
+        Returns:
+            E: total energy of shape ()
+        """
+        T = kinetic_energy_fn(params, q_d)
+        U = potential_energy_fn(params, q)
+        E = T + U
+
+        return E
+
+    auxiliary_fns = {
+        "apply_eps_to_bend_strains": apply_eps_to_bend_strains,
+        "kinetic_energy_fn": kinetic_energy_fn,
+        "potential_energy_fn": potential_energy_fn,
+        "energy_fn": energy_fn,
+    }
+
+    return B_xi, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns

tests/test_planar_pcs.py

@@ -27,7 +27,7 @@ def test_planar_pcs_one_segment():
     # activate all strains (i.e. bending, shear, and axial)
     strain_selector = jnp.ones((3,), dtype=bool)
 
-    strain_basis, forward_kinematics_fn, dynamical_matrices_fn = planar_pcs.factory(
+    strain_basis, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns = planar_pcs.factory(
         sym_exp_filepath, strain_selector
     )
     forward_dynamics_fn = partial(