Run formatting

Author: mstoelzle

examples/simulate_planar_hsa.py

@@ -59,7 +59,9 @@
         inverse_kinematics_end_effector_fn,
         dynamical_matrices_fn,
         sys_helpers,
-    ) = planar_hsa.factory(sym_exp_filepath, strain_selector, consider_hysteresis=consider_hysteresis)
+    ) = planar_hsa.factory(
+        sym_exp_filepath, strain_selector, consider_hysteresis=consider_hysteresis
+    )
 
     # import matplotlib.pyplot as plt
     # plt.plot(chi_ps[0, :], chi_ps[1, :])
@@ -91,7 +93,9 @@
         x0 = jnp.zeros((2 * q0.shape[0],))  # initial condition
     x0 = x0.at[: q0.shape[0]].set(q0)  # set initial configuration
 
-    ode_fn = planar_hsa.ode_factory(dynamical_matrices_fn, params, consider_hysteresis=consider_hysteresis)
+    ode_fn = planar_hsa.ode_factory(
+        dynamical_matrices_fn, params, consider_hysteresis=consider_hysteresis
+    )
     ode_term = ODETerm(ode_fn)
 
     sol = diffeqsolve(
@@ -118,7 +122,7 @@
         video_path,
         video_ts=video_ts,
         q_ts=sol.ys[:, :3],
-        video_width=video_width, 
-        video_height=video_height
+        video_width=video_width,
+        video_height=video_height,
     )
     print(f"Video saved at {video_path}")

src/jsrm/parameters/hsa_params.py

@@ -80,22 +80,34 @@ def generate_common_base_params(
         # mapping hysteresis displacements to the strains
         # per default, we only model the hysteresis on the axial strain
         B_z = jax.scipy.linalg.block_diag(
-            *[jnp.array([[0.0], [0.0], [1.0]]) for _ in range(num_segments)]  # assumes 3 strains per segment
-
+            *[
+                jnp.array([[0.0], [0.0], [1.0]]) for _ in range(num_segments)
+            ]  # assumes 3 strains per segment
         )
         params["hysteresis"]["basis"] = B_z
         # number of hysteresis states
         n_z = params["hysteresis"]["basis"].shape[1]
         # ratio of post-yield and pre-yield stiffness
         hys_alpha_val = 0.6
-        hys_alpha = jnp.array([(hys_alpha_val if B_z[xi_idx, :].sum() > 0 else 1.0) for xi_idx in range(3 * num_segments)])
+        hys_alpha = jnp.array(
+            [
+                (hys_alpha_val if B_z[xi_idx, :].sum() > 0 else 1.0)
+                for xi_idx in range(3 * num_segments)
+            ]
+        )
         params["hysteresis"]["alpha"] = hys_alpha
         # params["hysteresis"]["alpha"] = 0.6 * jnp.ones((n_q, ))  # ratio of post-yield and pre-yield stiffness
-        params["hysteresis"]["beta"] = 30.0 * jnp.ones((n_z, ))  # dimensionless parameter in the Bouc-Wen model
-        params["hysteresis"]["gamma"] = 1.0 * jnp.ones((n_z, ))  # dimensionless parameter in the Bouc-Wen model
-        params["hysteresis"]["n"] = 1.0 * jnp.ones((n_z, ))  # dimensionless parameter in the Bouc-Wen model
+        params["hysteresis"]["beta"] = 30.0 * jnp.ones(
+            (n_z,)
+        )  # dimensionless parameter in the Bouc-Wen model
+        params["hysteresis"]["gamma"] = 1.0 * jnp.ones(
+            (n_z,)
+        )  # dimensionless parameter in the Bouc-Wen model
+        params["hysteresis"]["n"] = 1.0 * jnp.ones(
+            (n_z,)
+        )  # dimensionless parameter in the Bouc-Wen model
         # to remove redundancy from Bouc-Wen model, choose A = 1
-        params["hysteresis"]["A"] = jnp.ones((n_z, ))
+        params["hysteresis"]["A"] = jnp.ones((n_z,))
 
     return params
 

src/jsrm/symbolic_derivation/planar_hsa.py

@@ -341,8 +341,10 @@ def symbolically_derive_planar_hsa_model(
             J_betar = sp.Matrix([[1, 0, 0], [0, 1, 0], [roff[i, j], 0, 1]])
 
             # nominal stiffness matrix of the virtual backbone
-            Shat[3 * i : 3 * (i + 1), 3 * i : 3 * (i + 1)] += J_betar.T @ Shatr @ J_betar
-            
+            Shat[3 * i : 3 * (i + 1), 3 * i : 3 * (i + 1)] += (
+                J_betar.T @ Shatr @ J_betar
+            )
+
             # contribution of elastic forces of current rod to the virtual backbone
             vKr = J_betar.T @ Shatr @ (pxir - pxi_eqr)
             # add contribution of elasticity vector

src/jsrm/systems/planar_hsa.py

@@ -150,9 +150,7 @@ def select_params_for_lambdify_fn(params: Dict[str, Array]) -> List[Array]:
     G_lambda = sp.lambdify(
         params_syms_cat + sym_exps["state_syms"]["xi"], sym_exps["exps"]["G"], "jax"
     )
-    Shat_lambda = sp.lambdify(
-        params_syms_cat, sym_exps["exps"]["Shat"], "jax"
-    )
+    Shat_lambda = sp.lambdify(params_syms_cat, sym_exps["exps"]["Shat"], "jax")
     K_lambda = sp.lambdify(
         params_syms_cat + sym_exps["state_syms"]["xi"], sym_exps["exps"]["K"], "jax"
     )
@@ -602,8 +600,8 @@ def dynamical_matrices_fn(
             B_z = params["hysteresis"]["basis"]
             hyst_alpha = params["hysteresis"]["alpha"]
             # add the post-yield potential forces
-            K = hyst_alpha*K + (1-hyst_alpha) * Shat @ (B_z @ z)
-            
+            K = hyst_alpha * K + (1 - hyst_alpha) * Shat @ (B_z @ z)
+
             # TODO: add post-yield potential forces (i.e., hysteresis effects) to the actuation vector
 
         # apply the strain basis
@@ -752,11 +750,16 @@ def ode_fn(t: float, x: Array, u: Array) -> Array:
 
             n_z = B_z.shape[1]
             n_q = (x.shape[0] - n_z) // 2
-            q, q_d, z = x[:n_q], x[n_q:2*n_q], x[2*n_q:]
+            q, q_d, z = x[:n_q], x[n_q : 2 * n_q], x[2 * n_q :]
 
             z_d = (B_z.T @ q_d) * (
-                    hys_params["A"] - jnp.abs(z)**hys_params["n"] * (hys_params["gamma"] + hys_params["beta"] * jnp.sign((B_z.T @ q_d) * z))
+                hys_params["A"]
+                - jnp.abs(z) ** hys_params["n"]
+                * (
+                    hys_params["gamma"]
+                    + hys_params["beta"] * jnp.sign((B_z.T @ q_d) * z)
                 )
+            )
         else:
             n_q = x.shape[0] // 2
             q, q_d = x[:n_q], x[n_q:]

src/jsrm/systems/utils.py

@@ -31,10 +31,10 @@ def substitute_params_into_all_symbolic_expressions(
     for exp_key, exp_val in exps.items():
         if issubclass(type(exp_val), list):
             for exp_item_idx, exp_item_val in enumerate(exp_val):
-                exps[exp_key][
-                    exp_item_idx
-                ] = substitute_params_into_single_symbolic_expression(
-                    exp_item_val, params_syms, params
+                exps[exp_key][exp_item_idx] = (
+                    substitute_params_into_single_symbolic_expression(
+                        exp_item_val, params_syms, params
+                    )
                 )
         else:
             exps[exp_key] = substitute_params_into_single_symbolic_expression(
@@ -74,7 +74,7 @@ def substitute_params_into_single_symbolic_expression(
 
 
 def concatenate_params_syms(
-    params_syms: Dict[str, Union[sp.Symbol, List[sp.Symbol]]]
+    params_syms: Dict[str, Union[sp.Symbol, List[sp.Symbol]]],
 ) -> List[sp.Symbol]:
     # concatenate the robot params symbols
     params_syms_cat = []