Correction test_planar_pcs to be consistent with

the class
Start of class for Planar HSA (to be continued)

Author: solangegbv

examples/simulate_class_planar_hsa.py

@@ -0,0 +1,279 @@
+import jax
+import jax.numpy as jnp
+from jax import jit, vmap
+
+import jsrm
+from jsrm.systems.class_planar_hsa import PlanarHSA
+from jsrm.parameters.hsa_params import PARAMS_FPU_CONTROL, PARAMS_FPU_HYSTERESIS_CONTROL
+
+from typing import Callable
+from jax import Array
+
+import numpy as onp
+
+from diffrax import Tsit5
+import cv2  # importing cv2
+
+from pathlib import Path
+
+jax.config.update("jax_enable_x64", True)  # double precision
+jnp.set_printoptions(
+    threshold=jnp.inf,
+    linewidth=jnp.inf,
+    formatter={"float_kind": lambda x: "0" if x == 0 else f"{x:.2e}"},
+)
+
+
+def draw_robot(
+    robot: PlanarHSA,
+    q: Array,
+    width: int = 700,
+    height: int = 700,
+    num_points: int = 50,
+) -> onp.ndarray:
+    """
+    Draw the robot in OpenCV.
+    Args:
+        robot: 
+        q: configuration as shape (3, )
+        width: image width
+        height: image height
+        num_points: number of points to plot along the length of the robot
+    """
+    # plotting in OpenCV
+    h, w = height, width  # img height and width
+    ppm = h / (
+        2.0 * jnp.sum(robot.params["lpc"] + robot.params["l"] + robot.params["ldc"])
+    )  # pixel per meter
+    base_color = (0, 0, 0)  # black base color in BGR
+    backbone_color = (255, 0, 0)  # blue robot color in BGR
+    rod_color = (0, 255, 0)  # green rod color in BGR
+    platform_color = (0, 0, 255)  # red platform color in BGR
+
+    batched_forward_kinematics_virtual_backbone_fn = vmap(
+        robot.forward_kinematics_virtual_backbone_fn,
+        in_axes=(None, 0), out_axes=-1
+    )
+    batched_forward_kinematics_rod_fn = vmap(
+        robot.forward_kinematics_rod_fn,
+        in_axes=(None, 0, None), out_axes=-1
+    )
+    batched_forward_kinematics_platform_fn = vmap(
+        robot.forward_kinematics_platform_fn,
+        in_axes=(None, 0), out_axes=0
+    )
+
+    L_max = jnp.sum(robot.params["l"])  # total length of the robot
+    # we use for plotting N points along the length of the robot
+    s_ps = jnp.linspace(0, L_max, num_points)
+
+    # poses along the robot of shape (3, N)
+    chiv_ps = batched_forward_kinematics_virtual_backbone_fn(q, s_ps)  # poses of virtual backbone
+    chiL_ps = batched_forward_kinematics_rod_fn(q, s_ps, 0)  # poses of left rod
+    chiR_ps = batched_forward_kinematics_rod_fn(q, s_ps, 1)  # poses of left rod
+    # poses of the platforms
+    chip_ps = batched_forward_kinematics_platform_fn(q, jnp.arange(0, robot.num_segments))
+
+    img = 255 * onp.ones((w, h, 3), dtype=jnp.uint8)  # initialize background to white
+    uv_robot_origin = onp.array(
+        [w // 2, h * (1 - 0.1)], dtype=jnp.int32
+    )  # in x-y pixel coordinates
+    uv_robot_origin_jax = jnp.array(uv_robot_origin)
+
+    @jit
+    def chi2u(chi: Array) -> Array:
+        """
+        Map Cartesian coordinates to pixel coordinates.
+        Args:
+            chi: Cartesian poses of shape (3)
+
+        Returns:
+            uv: pixel coordinates of shape (2)
+        """
+        uv_off = jnp.array((chi[1:] * ppm), dtype=jnp.int32)
+        # invert the v pixel coordinate
+        uv_off = uv_off.at[1].set(-uv_off[1])
+        # invert the v pixel coordinate
+        uv = uv_robot_origin_jax + uv_off
+        return uv
+
+    batched_chi2u = vmap(chi2u, in_axes=-1, out_axes=0)
+
+    # draw base
+    cv2.rectangle(img, (0, uv_robot_origin[1]), (w, h), color=base_color, thickness=-1)
+
+    # draw the virtual backbone
+    # add the first point of the proximal cap and the last point of the distal cap
+    chiv_ps = jnp.concatenate(
+        [
+            (chiv_ps[:, 0] - jnp.array([0.0, 0.0, params["lpc"][0]])).reshape(3, 1),
+            chiv_ps,
+            (
+                chiv_ps[:, -1]
+                + jnp.array(
+                    [
+                        chiv_ps[2, -1],
+                        -jnp.sin(chiv_ps[2, -1]) * params["ldc"][-1],
+                        jnp.cos(chiv_ps[2, -1]) * params["ldc"][-1],
+                    ]
+                )
+            ).reshape(3, 1),
+        ],
+        axis=1,
+    )
+    curve_virtual_backbone = onp.array(batched_chi2u(chiv_ps))
+    cv2.polylines(
+        img, [curve_virtual_backbone], isClosed=False, color=backbone_color, thickness=5
+    )
+
+    # draw the rods
+    # add the first point of the proximal cap and the last point of the distal cap
+    chiL_ps = jnp.concatenate(
+        [
+            (chiL_ps[:, 0] - jnp.array([0.0, 0.0, params["lpc"][0]])).reshape(3, 1),
+            chiL_ps,
+            (
+                chiL_ps[:, -1]
+                + jnp.array(
+                    [
+                        chiL_ps[2, -1],
+                        -jnp.sin(chiL_ps[2, -1]) * params["ldc"][-1],
+                        jnp.cos(chiL_ps[2, -1]) * params["ldc"][-1],
+                    ]
+                )
+            ).reshape(3, 1),
+        ],
+        axis=1,
+    )
+    curve_rod_left = onp.array(batched_chi2u(chiL_ps))
+    cv2.polylines(
+        img,
+        [curve_rod_left],
+        isClosed=False,
+        color=rod_color,
+        thickness=10,
+        # thickness=2*int(ppm * params["rout"].mean(axis=0)[0])
+    )
+    # add the first point of the proximal cap and the last point of the distal cap
+    chiR_ps = jnp.concatenate(
+        [
+            (chiR_ps[:, 0] - jnp.array([0.0, 0.0, params["lpc"][0]])).reshape(3, 1),
+            chiR_ps,
+            (
+                chiR_ps[:, -1]
+                + jnp.array(
+                    [
+                        chiR_ps[2, -1],
+                        -jnp.sin(chiR_ps[2, -1]) * params["ldc"][-1],
+                        jnp.cos(chiR_ps[2, -1]) * params["ldc"][-1],
+                    ]
+                )
+            ).reshape(3, 1),
+        ],
+        axis=1,
+    )
+    curve_rod_right = onp.array(batched_chi2u(chiR_ps))
+    cv2.polylines(img, [curve_rod_right], isClosed=False, color=rod_color, thickness=10)
+
+    # draw the platform
+    for i in range(chip_ps.shape[0]):
+        # iterate over the platforms
+        platform_R = jnp.array(
+            [
+                [jnp.cos(chip_ps[i, 0]), -jnp.sin(chip_ps[i, 0])],
+                [jnp.sin(chip_ps[i, 0]), jnp.cos(chip_ps[i, 0])],
+            ]
+        )  # rotation matrix for the platform
+        platform_llc = chip_ps[i, 1:] + platform_R @ jnp.array(
+            [
+                -params["pcudim"][i, 1] / 2,  # go half the width to the left
+                -params["pcudim"][i, 2] / 2,  # go half the height down
+            ]
+        )  # lower left corner of the platform
+        platform_ulc = chip_ps[i, 1:] + platform_R @ jnp.array(
+            [
+                -params["pcudim"][i, 1] / 2,  # go half the width to the left
+                +params["pcudim"][i, 2] / 2,  # go half the height down
+            ]
+        )  # upper left corner of the platform
+        platform_urc = chip_ps[i, 1:] + platform_R @ jnp.array(
+            [
+                +params["pcudim"][i, 1] / 2,  # go half the width to the left
+                +params["pcudim"][i, 2] / 2,  # go half the height down
+            ]
+        )  # upper right corner of the platform
+        platform_lrc = chip_ps[i, 1:] + platform_R @ jnp.array(
+            [
+                +params["pcudim"][i, 1] / 2,  # go half the width to the left
+                -params["pcudim"][i, 2] / 2,  # go half the height down
+            ]
+        )  # lower right corner of the platform
+        platform_curve = jnp.stack(
+            [platform_llc, platform_ulc, platform_urc, platform_lrc, platform_llc],
+            axis=1,
+        )
+        # cv2.polylines(img, [onp.array(batched_chi2u(platform_curve))], isClosed=True, color=platform_color, thickness=5)
+        cv2.fillPoly(
+            img, [onp.array(batched_chi2u(platform_curve))], color=platform_color
+        )
+
+    return img
+
+
+if __name__ == "__main__":
+    num_segments = 1
+    num_rods_per_segment = 2
+    
+    # filepath to symbolic expressions
+    sym_exp_filepath = (
+        Path(jsrm.__file__).parent
+        / "symbolic_expressions"
+        / f"planar_hsa_ns-{num_segments}_nrs-{num_rods_per_segment}.dill"
+    )
+    
+    # activate all strains (i.e. bending, shear, and axial)
+    strain_selector = jnp.ones((3 * num_segments,), dtype=bool)
+    consider_hysteresis = True
+    
+    params = PARAMS_FPU_HYSTERESIS_CONTROL if consider_hysteresis else PARAMS_FPU_CONTROL
+    # increase damping for simulation stability
+    params["zetab"] = 5 * params["zetab"]
+    params["zetash"] = 5 * params["zetash"]
+    params["zetaa"] = 5 * params["zetaa"]
+    
+    # ======================================================
+    # Robot initialization
+    # ======================================================
+    robot = PlanarHSA(
+        sym_exp_filepath=sym_exp_filepath,
+        params=params,
+        strain_selector=strain_selector,
+        consider_hysteresis=consider_hysteresis,
+    )
+    
+    # =====================================================
+    # Simulation upon time
+    # =====================================================
+    # Initial configuration
+    q0 = jnp.array([jnp.pi, 0.0, 0.0])
+    # Initial velocities
+    qd0 = jnp.zeros_like(q0)
+    # Motor actuation angles
+    phi = jnp.array([jnp.pi, jnp.pi / 2])
+    
+    # Displaying the image
+    window_name = f"Planar HSA with {num_segments} segments"
+    img = draw_robot(
+        robot,
+        q = q0,
+    )
+    
+    
+    # Simulation time parameters
+    t0 = 0.0
+    t1 = 5.0
+    dt = 5e-5  # time step
+    skip_step = 100  # how many time steps to skip in between video frames
+    
+    # Solver
+    solver = Tsit5()

examples/simulate_planar_pcs.py

@@ -3,7 +3,7 @@
 from jsrm.systems.planar_pcs import PlanarPCS
 import jax.numpy as jnp
 
-from typing import Callable, Dict
+from typing import Callable
 from jax import Array
 
 import numpy as onp