Run formatting

Author: mstoelzle

examples/analyze_segment_fusion.py

@@ -1,8 +1,12 @@
 import sympy as sp
 
-l1, l2 = sp.symbols('l1 l2', nonnegative=True, nonzero=True)
-kappa_be1, sigma_sh1, sigma_ax1 = sp.symbols('kappa_be1 sigma_sh1 sigma_ax1', real=True, nonnegative=True, nonzero=True)
-kappa_be2, sigma_sh2, sigma_ax2 = sp.symbols('kappa_be2 sigma_sh2 sigma_ax2', real=True, nonnegative=True, nonzero=True)
+l1, l2 = sp.symbols("l1 l2", nonnegative=True, nonzero=True)
+kappa_be1, sigma_sh1, sigma_ax1 = sp.symbols(
+    "kappa_be1 sigma_sh1 sigma_ax1", real=True, nonnegative=True, nonzero=True
+)
+kappa_be2, sigma_sh2, sigma_ax2 = sp.symbols(
+    "kappa_be2 sigma_sh2 sigma_ax2", real=True, nonnegative=True, nonzero=True
+)
 
 # deactivate shear and axial strains
 # sigma_sh1, sigma_sh2 = 0.0, 0.0
@@ -19,30 +23,42 @@
 # apply the forward kinematics on the individual segments
 s_be1, c_be1 = sp.sin(kappa_be1 * l1), sp.cos(kappa_be1 * l1)
 s_be2, c_be2 = sp.sin(kappa_be2 * l2), sp.cos(kappa_be2 * l2)
-chi1 = sp.Matrix([
-    [sigma_sh1 * s_be1 / kappa_be1 + sigma_ax1 * (c_be1 - 1) / kappa_be1],
-    [sigma_sh1 * (1 - c_be1) / kappa_be1 + sigma_ax1 * s_be1 / kappa_be1],
-    [kappa_be1 * l1],
-])
+chi1 = sp.Matrix(
+    [
+        [sigma_sh1 * s_be1 / kappa_be1 + sigma_ax1 * (c_be1 - 1) / kappa_be1],
+        [sigma_sh1 * (1 - c_be1) / kappa_be1 + sigma_ax1 * s_be1 / kappa_be1],
+        [kappa_be1 * l1],
+    ]
+)
 # rotation matrix
-R1 = sp.Matrix([
-    [c_be1, -s_be1, 0],
-    [s_be1, c_be1, 0],
-    [0, 0, 1],
-])
-chi2_rel = sp.Matrix([
-    [sigma_sh2 * s_be2 / kappa_be2 + sigma_ax2 * (c_be2 - 1) / kappa_be2],
-    [sigma_sh2 * (1 - c_be2) / kappa_be2 + sigma_ax2 * s_be2 / kappa_be2],
-    [kappa_be2 * l2],
-])
+R1 = sp.Matrix(
+    [
+        [c_be1, -s_be1, 0],
+        [s_be1, c_be1, 0],
+        [0, 0, 1],
+    ]
+)
+chi2_rel = sp.Matrix(
+    [
+        [sigma_sh2 * s_be2 / kappa_be2 + sigma_ax2 * (c_be2 - 1) / kappa_be2],
+        [sigma_sh2 * (1 - c_be2) / kappa_be2 + sigma_ax2 * s_be2 / kappa_be2],
+        [kappa_be2 * l2],
+    ]
+)
 chi2 = sp.simplify(chi1 + R1 @ chi2_rel)
 # print("chi2 =\n", chi2)
 # apply the one-segment inverse kinematic on the distal end of the two segments
 px, py, th = chi2[0, 0], chi2[1, 0], chi2[2, 0]
 s_th, c_th = sp.sin(th), sp.cos(th)
-qfkik = sp.simplify(th / (2 * (l1 + l2)) * sp.Matrix([
-    [2],
-    [py - px * s_th / (c_th - 1)],
-    [-px - py * s_th / (c_th - 1)],
-]))
+qfkik = sp.simplify(
+    th
+    / (2 * (l1 + l2))
+    * sp.Matrix(
+        [
+            [2],
+            [py - px * s_th / (c_th - 1)],
+            [-px - py * s_th / (c_th - 1)],
+        ]
+    )
+)
 print("qfkik =\n", qfkik)

examples/demo_planar_hsa_motor2ee_jacobian.py

@@ -26,7 +26,9 @@
 )
 
 
-def factory_fn(params: Dict[str, Array], verbose: bool = False) -> Tuple[Callable, Callable]:
+def factory_fn(
+    params: Dict[str, Array], verbose: bool = False
+) -> Tuple[Callable, Callable]:
     """
     Factory function for the planar HSA.
     Args:
@@ -45,7 +47,9 @@ def factory_fn(params: Dict[str, Array], verbose: bool = False) -> Tuple[Callabl
         sys_helpers,
     ) = planar_hsa.factory(sym_exp_filepath, strain_selector)
     dynamical_matrices_fn = partial(dynamical_matrices_fn, params)
-    forward_kinematics_end_effector_fn = jit(partial(forward_kinematics_end_effector_fn, params))
+    forward_kinematics_end_effector_fn = jit(
+        partial(forward_kinematics_end_effector_fn, params)
+    )
     jacobian_end_effector_fn = jit(partial(jacobian_end_effector_fn, params))
 
     def residual_fn(q: Array, phi: Array) -> Array:
@@ -64,7 +68,9 @@ def residual_fn(q: Array, phi: Array) -> Array:
     print("Compiling jac_residual_fn...")
     print(jac_residual_fn(jnp.zeros((3,)), jnp.zeros((2,))))
 
-    def phi2q_static_model_fn(phi: Array, q0: Array = jnp.zeros((3, ))) -> Tuple[Array, Dict[str, Array]]:
+    def phi2q_static_model_fn(
+        phi: Array, q0: Array = jnp.zeros((3,))
+    ) -> Tuple[Array, Dict[str, Array]]:
         """
         A static model mapping the motor angles to the planar HSA configuration using scipy.optimize.minimize.
         Arguments:
@@ -82,7 +88,12 @@ def phi2q_static_model_fn(phi: Array, q0: Array = jnp.zeros((3, ))) -> Tuple[Arr
             options={"disp": True} if verbose else None,
         )
         if verbose:
-         print("Optimization converged after", sol.nit, "iterations with residual", sol.fun)
+            print(
+                "Optimization converged after",
+                sol.nit,
+                "iterations with residual",
+                sol.fun,
+            )
 
         # configuration that minimizes the residual
         q = jnp.array(sol.x)
@@ -95,7 +106,9 @@ def phi2q_static_model_fn(phi: Array, q0: Array = jnp.zeros((3, ))) -> Tuple[Arr
 
         return q, aux
 
-    def phi2chi_static_model_fn(phi: Array, q0: Array = jnp.zeros((3, ))) -> Tuple[Array, Dict[str, Array]]:
+    def phi2chi_static_model_fn(
+        phi: Array, q0: Array = jnp.zeros((3,))
+    ) -> Tuple[Array, Dict[str, Array]]:
         """
         A static model mapping the motor angles to the planar end-effector pose.
         Arguments:
@@ -110,7 +123,9 @@ def phi2chi_static_model_fn(phi: Array, q0: Array = jnp.zeros((3, ))) -> Tuple[A
         aux["chi"] = chi
         return chi, aux
 
-    def jac_phi2chi_static_model_fn(phi: Array, q0: Array = jnp.zeros((3, ))) -> Tuple[Array, Dict[str, Array]]:
+    def jac_phi2chi_static_model_fn(
+        phi: Array, q0: Array = jnp.zeros((3,))
+    ) -> Tuple[Array, Dict[str, Array]]:
         """
         Compute the Jacobian between the actuation space and the task-space.
         Arguments:
@@ -157,17 +172,12 @@ def jac_phi2chi_static_model_fn(phi: Array, q0: Array = jnp.zeros((3, ))) -> Tup
                 phi = jnp.array([1.0, 1.0])
             case _:
                 rng, subkey = random.split(rng)
-                phi = random.uniform(
-                    subkey,
-                    phi_max.shape,
-                    minval=0.0,
-                    maxval=phi_max
-                )
+                phi = random.uniform(subkey, phi_max.shape, minval=0.0, maxval=phi_max)
 
         print("i", i)
 
         chi, aux = phi2chi_static_model_fn(phi, q0=q0)
         print("phi", phi, "q", aux["q"], "chi", chi)
 
         J_phi2chi, aux = jac_phi2chi_static_model_fn(phi, q0=q0)
-        print("J_phi2chi:\n", J_phi2chi)
+        print("J_phi2chi:\n", J_phi2chi)

examples/derive_planar_pcs.py

@@ -12,5 +12,7 @@
         / f"planar_pcs_ns-{NUM_SEGMENTS}.dill"
     )
     symbolically_derive_planar_pcs_model(
-        num_segments=NUM_SEGMENTS, filepath=sym_exp_filepath, simplify_expressions=True if NUM_SEGMENTS < 3 else False
+        num_segments=NUM_SEGMENTS,
+        filepath=sym_exp_filepath,
+        simplify_expressions=True if NUM_SEGMENTS < 3 else False,
     )

examples/simulate_planar_pcs.py

@@ -95,8 +95,8 @@ def draw_robot(
 
 
 if __name__ == "__main__":
-    strain_basis, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns = planar_pcs.factory(
-        sym_exp_filepath, strain_selector
+    strain_basis, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns = (
+        planar_pcs.factory(sym_exp_filepath, strain_selector)
     )
     batched_forward_kinematics = vmap(
         forward_kinematics_fn, in_axes=(None, None, 0), out_axes=-1
@@ -143,7 +143,9 @@ def draw_robot(
     q_d_ts = sol.ys[:, n_q:]
 
     # evaluate the forward kinematics along the trajectory
-    chi_ee_ts = vmap(forward_kinematics_fn, in_axes=(None, 0, None))(params, q_ts, jnp.array([jnp.sum(params["l"])]))
+    chi_ee_ts = vmap(forward_kinematics_fn, in_axes=(None, 0, None))(
+        params, q_ts, jnp.array([jnp.sum(params["l"])])
+    )
     # plot the end-effector position along the trajectory
     plt.figure()
     plt.plot(chi_ee_ts[0, :], chi_ee_ts[1, :])
@@ -166,8 +168,12 @@ def draw_robot(
     plt.show()
 
     # 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))
+    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()

src/jsrm/symbolic_derivation/planar_hsa.py

@@ -507,7 +507,7 @@ def symbolically_derive_planar_hsa_model(
         Bpl = sp.simplify(Bpl)
     B = B + Bpl
     # add the gravitational potential energy of the payload
-    Upl = sp.simplify(- mpl * g.T @ pCoGpl)
+    Upl = sp.simplify(-mpl * g.T @ pCoGpl)
     U_g = U_g + Upl
 
     # simplify mass matrix

src/jsrm/symbolic_derivation/planar_pcs.py

@@ -7,7 +7,9 @@
 
 
 def symbolically_derive_planar_pcs_model(
-    num_segments: int, filepath: Optional[Union[str, Path]] = None, simplify_expressions: bool = True
+    num_segments: int,
+    filepath: Optional[Union[str, Path]] = None,
+    simplify_expressions: bool = True,
 ) -> Dict:
     """
     Symbolically derive the kinematics and dynamics of a planar continuum soft robot modelled with
@@ -184,9 +186,13 @@ def symbolically_derive_planar_pcs_model(
                 s, l[-1]
             ),  # expression for end-effector pose of shape (3, )
             "J_sms": J_sms,  # list of Jacobians (for each segment) of shape (3, num_dof)
-            "Jee": J_sms[-1].subs(s, l[-1]),  # end-effector Jacobian of shape (3, num_dof)
+            "Jee": J_sms[-1].subs(
+                s, l[-1]
+            ),  # end-effector Jacobian of shape (3, num_dof)
             "J_d_sms": J_d_sms,  # list of time derivatives of Jacobians (for each segment)
-            "Jee_d": J_d_sms[-1].subs(s, l[-1]),  # time derivative of end-effector Jacobian of shape (3, num_dof)
+            "Jee_d": J_d_sms[-1].subs(
+                s, l[-1]
+            ),  # time derivative of end-effector Jacobian of shape (3, num_dof)
             "B": B,  # mass matrix
             "C": C,  # coriolis matrix
             "G": G,  # gravity vector

src/jsrm/systems/planar_pcs.py

@@ -175,7 +175,7 @@ def apply_eps_to_bend_strains(xi: Array, _eps: float) -> Array:
 
         return xi_epsed
 
-    def classify_segment(params: Dict[str, Array], s: Array) -> Tuple[Array, Array] :
+    def classify_segment(params: Dict[str, Array], s: Array) -> Tuple[Array, Array]:
         """
         Classify the point along the robot to the corresponding segment
         Args:
@@ -192,14 +192,16 @@ def classify_segment(params: Dict[str, Array], s: Array) -> Tuple[Array, Array]
         # determine in which segment the point is located
         # use argmax to find the last index where the condition is true
         segment_idx = (
-                l_cum.shape[0] - 1 - jnp.argmax((s >= l_cum_padded[:-1])[::-1]).astype(int)
+            l_cum.shape[0] - 1 - jnp.argmax((s >= l_cum_padded[:-1])[::-1]).astype(int)
         )
         # point coordinate along the segment in the interval [0, l_segment]
         s_segment = s - l_cum_padded[segment_idx]
 
         return segment_idx, s_segment
 
-    def stiffness_fn(params: Dict[str, Array], formulate_in_strain_space: bool = False) -> Array:
+    def stiffness_fn(
+        params: Dict[str, Array], formulate_in_strain_space: bool = False
+    ) -> Array:
         """
         Compute the stiffness matrix of the system.
         Args:
@@ -210,7 +212,7 @@ def stiffness_fn(params: Dict[str, Array], formulate_in_strain_space: bool = Fal
         """
         # cross-sectional area and second moment of area
         A = jnp.pi * params["r"] ** 2
-        Ib = A ** 2 / (4 * jnp.pi)
+        Ib = A**2 / (4 * jnp.pi)
 
         # elastic and shear modulus
         E, G = params["E"], params["G"]
@@ -260,7 +262,7 @@ def forward_kinematics_fn(
 
     @jit
     def jacobian_fn(
-            params: Dict[str, Array], q: Array, s: Array, eps: float = global_eps
+        params: Dict[str, Array], q: Array, s: Array, eps: float = global_eps
     ) -> Array:
         """
         Evaluate the forward kinematics the tip of the links
@@ -341,7 +343,6 @@ def dynamical_matrices_fn(
         alpha = B_xi.T @ jnp.identity(n_xi) @ B_xi
 
         return B, C, G, K, D, alpha
-    
 
     def kinetic_energy_fn(params: Dict[str, Array], q: Array, q_d: Array) -> Array:
         """
@@ -357,11 +358,12 @@ def kinetic_energy_fn(params: Dict[str, Array], q: Array, q_d: Array) -> Array:
 
         # 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:
+    def potential_energy_fn(
+        params: Dict[str, Array], q: Array, eps: float = 1e4 * global_eps
+    ) -> Array:
         """
         Compute the potential energy of the system.
         Args:
@@ -460,7 +462,12 @@ def operational_space_dynamical_matrices_fn(
             segment_idx, J_lambda_sms, *params_for_lambdify, *xi_epsed, s_segment
         ).squeeze()
         J_d = lax.switch(
-            segment_idx, J_d_lambda_sms, *params_for_lambdify, *xi_epsed, *xi_d, s_segment
+            segment_idx,
+            J_d_lambda_sms,
+            *params_for_lambdify,
+            *xi_epsed,
+            *xi_d,
+            s_segment,
         ).squeeze()
         # apply the operational_space_selector and strain basis to the J and J_d
         J = J[operational_space_selector, :] @ B_xi