Properly implement actuation_basis

examples/simulate_pneumatic_planar_pcs.py

@@ -47,6 +47,18 @@
 # strain_selector = jnp.ones((3 * num_segments,), dtype=bool)
 strain_selector = jnp.array([True, False, True])[None, :].repeat(num_segments, axis=0).flatten()
 
+B_xi, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns = (
+    pneumatic_planar_pcs.factory(num_segments, sym_exp_filepath, strain_selector)
+)
+# jit the functions
+dynamical_matrices_fn = jax.jit(dynamical_matrices_fn)
+actuation_mapping_fn = auxiliary_fns["actuation_mapping_fn"]
+
+def sweep_actuation_mapping():
+    q = jnp.zeros((2 * num_segments,))
+    A = actuation_mapping_fn(params, B_xi, q)
+    print("A =\n", A)
+
 
 def simulate_robot():
     # define initial configuration
@@ -59,12 +71,6 @@ def simulate_robot():
     sim_dt = 5e-5  # simulation time step
     ts = jnp.arange(0.0, 2, dt)  # time steps
 
-    strain_basis, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns = (
-        pneumatic_planar_pcs.factory(sym_exp_filepath, strain_selector)
-    )
-    # jit the functions
-    dynamical_matrices_fn = jax.jit(partial(dynamical_matrices_fn))
-
     x0 = jnp.concatenate([q0, jnp.zeros_like(q0)])  # initial condition
     tau = jnp.zeros_like(q0)  # torques
 
@@ -160,4 +166,5 @@ def simulate_robot():
     plt.show()
 
 if __name__ == "__main__":
+    sweep_actuation_mapping()
     simulate_robot()

src/jsrm/systems/pneumatic_planar_pcs.py

@@ -3,7 +3,7 @@
 import jax.numpy as jnp
 from jsrm.math_utils import blk_diag
 import numpy as onp
-from typing import Dict, Tuple, Union
+from typing import Dict, Optional, Tuple, Union
 
 from .planar_pcs import factory as planar_pcs_factory
 
@@ -16,6 +16,7 @@ def factory(
     """
     Factory function for the planar PCS.
     Args:
+        num_segments: number of segments
         segment_actuation_selector: actuation selector for the segments as boolean array of shape (num_segments,)
             True entries signify that the segment is actuated, False entries signify that the segment is passive
     Returns:
@@ -27,7 +28,13 @@ def factory(
     actuation_dim = segment_actuation_selector.sum() * 2
 
     # matrix that maps the (possibly) underactuated actuation space to a full actuation space
-    actuation_basis = jnp.eye(num_segments)[segment_actuation_selector]
+    actuation_basis = jnp.zeros((2 * num_segments, actuation_dim))
+    actuation_basis_cumsum = jnp.cumsum(segment_actuation_selector)
+    for i in range(num_segments):
+        j = int(actuation_basis_cumsum[i].item()) - 1
+        if segment_actuation_selector[i].item() is True:
+            actuation_basis = actuation_basis.at[2 * i, j].set(1.0)
+            actuation_basis = actuation_basis.at[2 * i + 1, j + 1].set(1.0)
 
     def actuation_mapping_fn(
         params: Dict[str, Array],