Add plot backup, full-JAX PCS planar + Gaussian quadrature integratio…

src/jsrm/math_utils.py

@@ -1,9 +1,10 @@
 from jax import numpy as jnp
 from jax import Array, lax, jit
 
-
 @jit
-def blk_diag(a: Array) -> Array:
+def blk_diag(
+    a: Array
+) -> Array:
     """
     Create a block diagonal matrix from a tensor of blocks
     Args:
@@ -31,7 +32,7 @@ def assign_block_diagonal(i, _b):
 
     # Implement for loop to assign each block in `a` to the block-diagonal of `b`
     # Hint: use `jax.lax.fori_loop` and pass `assign_block_diagonal` as an argument
-    b = jnp.zeros((a.shape[0] * a.shape[1], a.shape[0] * a.shape[2]))
+    b = jnp.zeros((a.shape[0] * a.shape[1], a.shape[0] * a.shape[2]), dtype=a.dtype)
     b = lax.fori_loop(
         lower=0,
         upper=a.shape[0],
@@ -40,3 +41,33 @@ def assign_block_diagonal(i, _b):
     )
 
     return b
+
+@jit
+def blk_concat(
+    a: Array
+) -> Array:
+    """
+    Concatenate the matrices along the first axis
+
+    Args:
+        a (Array): matrices to be concatenated of shape (N, a, b)
+
+    Returns:
+        Array: concatenated matrix of shape (a, b * N)
+    """
+    b = a.transpose(1, 0, 2).reshape(a.shape[1], -1)
+    return b
+
+if __name__ == "__main__":
+    # Example usage
+    a = jnp.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
+    print("Original array:")
+    print(a)
+
+    b = blk_diag(a)
+    print("Block diagonal matrix:")
+    print(b)
+
+    c = blk_concat(a)
+    print("Concatenated matrix:")
+    print(c)

src/jsrm/symbolic_derivation/planar_pcs.py

@@ -60,9 +60,9 @@ def symbolically_derive_planar_pcs_model(
     # Jacobians (positional + orientation) in each segment as a function of the point coordinate s and its time derivative
     J_sms, J_d_sms = [], []
     # cross-sectional area of each segment
-    A = sp.zeros(num_segments)
+    A = sp.zeros(num_segments, 1)
     # second area moment of inertia of each segment
-    I = sp.zeros(num_segments)
+    I = sp.zeros(num_segments, 1)
     # inertia matrix
     B = sp.zeros(num_dof, num_dof)
     # potential energy

src/jsrm/systems/planar_pcs.py

@@ -472,7 +472,7 @@ def operational_space_dynamical_matrices_fn(
             Lambda: inertia matrix in the operational space of shape (3, 3)
             mu: matrix with corioli and centrifugal terms in the operational space of shape (3, 3)
             J: Jacobian of the Cartesian pose with respect to the generalized coordinates of shape (3, n_q)
-            J: time-derivative of the Jacobian of the end-effector pose with respect to the generalized coordinates
+            J_d: time-derivative of the Jacobian of the end-effector pose with respect to the generalized coordinates
                 of shape (3, n_q)
             JB_pinv: Dynamically-consistent pseudo-inverse of the Jacobian. Allows the mapping of torques
                 from the generalized coordinates to the operational space: f = JB_pinv.T @ tau_q