removed snake_id options and using native MuscularTorques for lateral…

… wave

examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_with_lifting_wave.py

@@ -14,7 +14,7 @@
     plot_curvature,
 )
 from snake_forcing import (
-    MuscleTorques_snake,
+    MuscleTorques_lifting,
 )
 from snake_contact import SnakeRodPlaneContact
 
@@ -31,36 +31,24 @@ def run_snake(
     SAVE_FIGURE=False,
     SAVE_VIDEO=False,
     SAVE_RESULTS=False,
-    args=[],
 ):
     # Initialize the simulation class
     snake_sim = SnakeSimulator()
 
-    # Get parser args
-    phase_space_params = args
-
     # Simulation parameters
     period = 2.0
     final_time = 20.0
-    time_step = phase_space_params.timestep
+    time_step = 5e-5
     total_steps = int(final_time / time_step)
     rendering_fps = 100
     step_skip = int(1.0 / (rendering_fps * time_step))
 
     # collection of snake characteristics
-    n_elem_collect = np.array([25, 50])
-    base_length_collect = np.array([0.35, 0.8])
-    base_radius_collect = np.array([0.009, 0.009])
-    snake_torque_ratio_collect = np.array([30.0, 20.0])
-    snake_torque_lift_ratio_collect = np.array([10.0, 20.0])
-
-    # select snake to run
-    snake_ID = 0
-
-    # setting up test params
-    n_elem = n_elem_collect[snake_ID]
-    base_length = base_length_collect[snake_ID]
-    base_radius = base_radius_collect[snake_ID]
+    n_elem = 25
+    base_length = 0.35
+    base_radius = 0.009
+    snake_torque_ratio = 30.0
+    snake_torque_liftratio = 10.0
 
     start = np.array([0.0, 0.0, 0.0 + base_radius])
     direction = np.array([1.0, 0.0, 0.0])
@@ -83,7 +71,7 @@ def run_snake(
     )
 
     snake_sim.append(shearable_rod)
-    damping_constant = phase_space_params.damping
+    damping_constant = 1e-1
 
     # use linear damping with constant damping ratio
     snake_sim.dampen(shearable_rod).using(
@@ -101,14 +89,12 @@ def run_snake(
 
     # 1. Add muscle torques -- lateral wave
     # Define lateral wave parameters
-    lateral_wave_length = phase_space_params.wave_length
-    snake_torque_ratio = snake_torque_ratio_collect[snake_ID]
-    snake_torque_liftratio = snake_torque_lift_ratio_collect[snake_ID]
+    lateral_wave_length = 1.0
     lateral_amp = b_coeff_lat[:-1]
 
     lateral_ratio = 1.0  # switch of lateral wave
     snake_sim.add_forcing_to(shearable_rod).using(
-        MuscleTorques_snake,
+        MuscleTorques,
         base_length=base_length,
         b_coeff=snake_torque_ratio * lateral_ratio * lateral_amp,
         period=period,
@@ -128,7 +114,7 @@ def run_snake(
     lift_ratio = 1.0  # switch of lifting wave
     phase = 0.5
     snake_sim.add_forcing_to(shearable_rod).using(
-        MuscleTorques_snake,
+        MuscleTorques_lifting,
         base_length=base_length,
         b_coeff=snake_torque_liftratio * lift_ratio * lift_amp,
         period=period,
@@ -236,28 +222,6 @@ def make_callback(self, system, time, current_step: int):
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument(
-        "--wave_length",
-        type=float,
-        default=1.0,
-    )
-    parser.add_argument(
-        "--timestep",
-        type=float,
-        default=5e-5,
-    )
-    parser.add_argument(
-        "--damping",
-        type=float,
-        default=1e-1,
-    )
-    args = parser.parse_args()
-
-    # print(args.wave_length)
-    # print(args.block_size)
-
     # Options
     PLOT_FIGURE = True
     SAVE_FIGURE = False
@@ -307,7 +271,7 @@ def optimize_snake(spline_coefficient):
 
         # run the simulation
         [avg_forward, avg_lateral, pp_list] = run_snake(
-            t_coeff_optimized, PLOT_FIGURE, SAVE_FIGURE, SAVE_VIDEO, SAVE_RESULTS, args
+            t_coeff_optimized, PLOT_FIGURE, SAVE_FIGURE, SAVE_VIDEO, SAVE_RESULTS
         )
 
         print("average forward velocity:", avg_forward)

examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py

@@ -17,7 +17,7 @@
 )
 
 
-class MuscleTorques_snake(NoForces):
+class MuscleTorques_lifting(NoForces):
     """
     This class applies muscle torques along the body. The applied muscle torques are treated
     as applied external forces. This class can apply
@@ -92,7 +92,7 @@ def __init__(
                 check if it is lateral muscle torque.
 
         """
-        super(MuscleTorques_snake, self).__init__()
+        super(MuscleTorques_lifting, self).__init__()
 
         self.direction = direction  # Direction torque applied
         self.angular_frequency = 2.0 * np.pi / period
@@ -189,24 +189,21 @@ def compute_muscle_torques(
             )
             # Head and tail of the snake is opposite compared to elastica cpp. We need to iterate torque_mag
             # from last to first element.
-            if is_lateral_wave:
-                torque = _batch_product_i_k_to_ik(direction, torque_mag[-2::-1])
-            else:
-                # compute torque direction for lifting wave.
-                # Here, direction of each element is computed separately
-                # based on the rod tangent and normal direction. This is implemented to
-                # correct the binormal direction when snake undergoes lateral bending
-                avg_element_direction = 0.5 * (tangents[..., :-1] + tangents[..., 1:])
-                torque_direction = _batch_vec_oneD_vec_cross(
-                    avg_element_direction, direction
-                )
-                torque_direction_unit = _batch_product_k_ik_to_ik(
-                    1 / (_batch_norm(torque_direction) + 1e-14),
-                    torque_direction,
-                )
-                torque = _batch_product_k_ik_to_ik(
-                    torque_mag[-2::-1], torque_direction_unit
-                )
+            # compute torque direction for lifting wave.
+            # Here, direction of each element is computed separately
+            # based on the rod tangent and normal direction. This is implemented to
+            # correct the binormal direction when snake undergoes lateral bending
+            avg_element_direction = 0.5 * (tangents[..., :-1] + tangents[..., 1:])
+            torque_direction = _batch_vec_oneD_vec_cross(
+                avg_element_direction, direction
+            )
+            torque_direction_unit = _batch_product_k_ik_to_ik(
+                1 / (_batch_norm(torque_direction) + 1e-14),
+                torque_direction,
+            )
+            torque = _batch_product_k_ik_to_ik(
+                torque_mag[-2::-1], torque_direction_unit
+            )
 
             inplace_addition(
                 external_torques[..., 1:],