Add B_xi to stiffness_fn interface

Author: mstoelzle

examples/simulate_planar_pcs.py

@@ -162,7 +162,7 @@ def draw_robot(
         )
         plt.plot(
             video_ts, q_ts[:, 3 * segment_idx + 2],
-            label=r"$\sigma_\mathrm{el," + str(segment_idx + 1) + "}$ [-]"
+            label=r"$\sigma_\mathrm{ax," + str(segment_idx + 1) + "}$ [-]"
         )
     plt.xlabel("Time [s]")
     plt.ylabel("Configuration")

src/jsrm/systems/planar_pcs.py

@@ -37,7 +37,7 @@ def factory(
                 strain are active / non-zero
         xi_eq: array of shape (3 * num_segments) with the rest strains of the rod
         stiffness_fn: function to compute the stiffness matrix of the system. Should have the signature
-            stiffness_fn(params: Dict[str, Array], formulate_in_strain_space: bool) -> Array
+            stiffness_fn(params: Dict[str, Array], B_xi, formulate_in_strain_space: bool) -> Array
         global_eps: small number to avoid singularities (e.g., division by zero)
     Returns:
         B_xi: strain basis matrix of shape (3 * num_segments, n_q)
@@ -204,12 +204,13 @@ def classify_segment(params: Dict[str, Array], s: Array) -> Tuple[Array, Array]:
 
     if stiffness_fn is None:
         def stiffness_fn(
-            params: Dict[str, Array], formulate_in_strain_space: bool = False
+            params: Dict[str, Array], B_xi: Array, formulate_in_strain_space: bool = False
         ) -> Array:
             """
             Compute the stiffness matrix of the system.
             Args:
                 params: Dictionary of robot parameters
+                B_xi: Strain basis matrix
                 formulate_in_strain_space: whether to formulate the elastic matrix in the strain space
             Returns:
                 K: elastic matrix of shape (n_q, n_q) if formulate_in_strain_space is False or (n_xi, n_xi) otherwise
@@ -329,7 +330,7 @@ def dynamical_matrices_fn(
         xi_epsed = apply_eps_to_bend_strains(xi, eps)
 
         # compute the stiffness matrix
-        K = stiffness_fn(params, formulate_in_strain_space=True)
+        K = stiffness_fn(params, B_xi, formulate_in_strain_space=True)
 
         # dissipative matrix from the parameters
         D = params.get("D", jnp.zeros((n_xi, n_xi)))
@@ -385,7 +386,7 @@ def potential_energy_fn(
         xi_epsed = apply_eps_to_bend_strains(xi, eps)
 
         # compute the stiffness matrix
-        K = stiffness_fn(params, formulate_in_strain_space=True)
+        K = stiffness_fn(params, B_xi, formulate_in_strain_space=True)
         # elastic energy
         U_K = (xi - xi_eq).T @ K @ (xi - xi_eq)  # evaluate K(xi) = K @ xi
 

src/jsrm/systems/utils.py

@@ -121,6 +121,6 @@ def compute_planar_stiffness_matrix(l: Array, A: Array, Ib: Array, E: Array, G:
     Returns:
         S: stiffness matrix of shape (3, 3)
     """
-    S = l* jnp.diag(jnp.stack([Ib * E, 4 / 3 * A * G, A * E], axis=0))
+    S = l * jnp.diag(jnp.stack([Ib * E, 4 / 3 * A * G, A * E], axis=0))
 
     return S