From 81c5b0a51c66e62e1083e5ed45f4bc3ccaa57d3e Mon Sep 17 00:00:00 2001
From: Ali-7800 <47090295+Ali-7800@users.noreply.github.com>
Date: Tue, 14 Nov 2023 20:01:43 -0600
Subject: [PATCH 1/9] Refactor: AnisotropicFrictionalPlane ->
 RodPlaneContactWithAnisotropicFriction

---
 .../ContinuumSnakeCase/continuum_snake.py     | 19 ++++++++++++-------
 .../FrictionValidationCases/axial_friction.py | 12 +++++++-----
 .../rolling_friction_initial_velocity.py      | 12 +++++++-----
 .../rolling_friction_on_inclined_plane.py     | 12 +++++++-----
 .../rolling_friction_torque.py                | 13 ++++++++-----
 examples/MuscularSnake/muscular_snake.py      | 10 ++++++----
 6 files changed, 47 insertions(+), 31 deletions(-)

diff --git a/examples/ContinuumSnakeCase/continuum_snake.py b/examples/ContinuumSnakeCase/continuum_snake.py
index e253f041..8987b155 100644
--- a/examples/ContinuumSnakeCase/continuum_snake.py
+++ b/examples/ContinuumSnakeCase/continuum_snake.py
@@ -13,7 +13,12 @@
 
 
 class SnakeSimulator(
-    ea.BaseSystemCollection, ea.Constraints, ea.Forcing, ea.Damping, ea.CallBacks
+    ea.BaseSystemCollection,
+    ea.Constraints,
+    ea.Forcing,
+    ea.Damping,
+    ea.CallBacks,
+    ea.Contact,
 ):
     pass
 
@@ -76,8 +81,10 @@ def run_snake(
     )
 
     # Add friction forces
-    origin_plane = np.array([0.0, -base_radius, 0.0])
-    normal_plane = normal
+    ground_plane = ea.Plane(
+        plane_origin=np.array([0.0, -base_radius, 0.0]), plane_normal=normal
+    )
+    snake_sim.append(ground_plane)
     slip_velocity_tol = 1e-8
     froude = 0.1
     mu = base_length / (period * period * np.abs(gravitational_acc) * froude)
@@ -85,12 +92,10 @@ def run_snake(
         [mu, 1.5 * mu, 2.0 * mu]
     )  # [forward, backward, sideways]
     static_mu_array = np.zeros(kinetic_mu_array.shape)
-    snake_sim.add_forcing_to(shearable_rod).using(
-        ea.AnisotropicFrictionalPlane,
+    snake_sim.detect_contact_between(shearable_rod, ground_plane).using(
+        ea.RodPlaneContactWithAnisotropicFriction,
         k=1.0,
         nu=1e-6,
-        plane_origin=origin_plane,
-        plane_normal=normal_plane,
         slip_velocity_tol=slip_velocity_tol,
         static_mu_array=static_mu_array,
         kinetic_mu_array=kinetic_mu_array,
diff --git a/examples/FrictionValidationCases/axial_friction.py b/examples/FrictionValidationCases/axial_friction.py
index 9c7d654d..d500a1e6 100644
--- a/examples/FrictionValidationCases/axial_friction.py
+++ b/examples/FrictionValidationCases/axial_friction.py
@@ -7,7 +7,9 @@
 )
 
 
-class AxialFrictionSimulator(ea.BaseSystemCollection, ea.Constraints, ea.Forcing):
+class AxialFrictionSimulator(
+    ea.BaseSystemCollection, ea.Constraints, ea.Forcing, ea.Contact
+):
     pass
 
 
@@ -74,13 +76,13 @@ def simulate_axial_friction_with(force=0.0):
     slip_velocity_tol = 1e-4
     static_mu_array = np.array([0.8, 0.4, 0.4])  # [forward, backward, sideways]
     kinetic_mu_array = np.array([0.4, 0.2, 0.2])  # [forward, backward, sideways]
+    friction_plane = ea.Plane(plane_origin=origin_plane, plane_normal=normal_plane)
+    axial_friction_sim.append(friction_plane)
 
-    axial_friction_sim.add_forcing_to(shearable_rod).using(
-        ea.AnisotropicFrictionalPlane,
+    axial_friction_sim.detect_contact_between(shearable_rod, friction_plane).using(
+        ea.RodPlaneContactWithAnisotropicFriction,
         k=10.0,
         nu=1e-4,
-        plane_origin=origin_plane,
-        plane_normal=normal_plane,
         slip_velocity_tol=slip_velocity_tol,
         static_mu_array=static_mu_array,
         kinetic_mu_array=kinetic_mu_array,
diff --git a/examples/FrictionValidationCases/rolling_friction_initial_velocity.py b/examples/FrictionValidationCases/rolling_friction_initial_velocity.py
index c74e9b85..c42edd7a 100644
--- a/examples/FrictionValidationCases/rolling_friction_initial_velocity.py
+++ b/examples/FrictionValidationCases/rolling_friction_initial_velocity.py
@@ -9,7 +9,7 @@
 
 
 class RollingFrictionInitialVelocitySimulator(
-    ea.BaseSystemCollection, ea.Constraints, ea.Forcing, ea.Damping
+    ea.BaseSystemCollection, ea.Constraints, ea.Forcing, ea.Damping, ea.Contact
 ):
     pass
 
@@ -88,13 +88,15 @@ def simulate_rolling_friction_initial_velocity_with(IFactor=0.0):
     slip_velocity_tol = 1e-6
     static_mu_array = np.array([0.4, 0.4, 0.4])  # [forward, backward, sideways]
     kinetic_mu_array = np.array([0.2, 0.2, 0.2])  # [forward, backward, sideways]
+    friction_plane = ea.Plane(plane_origin=origin_plane, plane_normal=normal_plane)
+    rolling_friction_initial_velocity_sim.append(friction_plane)
 
-    rolling_friction_initial_velocity_sim.add_forcing_to(shearable_rod).using(
-        ea.AnisotropicFrictionalPlane,
+    rolling_friction_initial_velocity_sim.detect_contact_between(
+        shearable_rod, friction_plane
+    ).using(
+        ea.RodPlaneContactWithAnisotropicFriction,
         k=10.0,
         nu=1e-4,
-        plane_origin=origin_plane,
-        plane_normal=normal_plane,
         slip_velocity_tol=slip_velocity_tol,
         static_mu_array=static_mu_array,
         kinetic_mu_array=kinetic_mu_array,
diff --git a/examples/FrictionValidationCases/rolling_friction_on_inclined_plane.py b/examples/FrictionValidationCases/rolling_friction_on_inclined_plane.py
index 57b04cd9..6685b252 100644
--- a/examples/FrictionValidationCases/rolling_friction_on_inclined_plane.py
+++ b/examples/FrictionValidationCases/rolling_friction_on_inclined_plane.py
@@ -9,7 +9,7 @@
 
 
 class RollingFrictionOnInclinedPlaneSimulator(
-    ea.BaseSystemCollection, ea.Constraints, ea.Forcing
+    ea.BaseSystemCollection, ea.Constraints, ea.Forcing, ea.Contact
 ):
     pass
 
@@ -74,13 +74,15 @@ def simulate_rolling_friction_on_inclined_plane_with(alpha_s=0.0):
     slip_velocity_tol = 1e-4
     static_mu_array = np.array([0.4, 0.4, 0.4])  # [forward, backward, sideways]
     kinetic_mu_array = np.array([0.2, 0.2, 0.2])  # [forward, backward, sideways]
+    friction_plane = ea.Plane(plane_origin=origin_plane, plane_normal=normal_plane)
+    rolling_friction_on_inclined_plane_sim.append(friction_plane)
 
-    rolling_friction_on_inclined_plane_sim.add_forcing_to(shearable_rod).using(
-        ea.AnisotropicFrictionalPlane,
+    rolling_friction_on_inclined_plane_sim.detect_contact_between(
+        shearable_rod, friction_plane
+    ).using(
+        ea.RodPlaneContactWithAnisotropicFriction,
         k=10.0,
         nu=1e-4,
-        plane_origin=origin_plane,
-        plane_normal=normal_plane,
         slip_velocity_tol=slip_velocity_tol,
         static_mu_array=static_mu_array,
         kinetic_mu_array=kinetic_mu_array,
diff --git a/examples/FrictionValidationCases/rolling_friction_torque.py b/examples/FrictionValidationCases/rolling_friction_torque.py
index 350a6cde..f4a0c22f 100644
--- a/examples/FrictionValidationCases/rolling_friction_torque.py
+++ b/examples/FrictionValidationCases/rolling_friction_torque.py
@@ -9,7 +9,7 @@
 
 
 class RollingFrictionTorqueSimulator(
-    ea.BaseSystemCollection, ea.Constraints, ea.Forcing
+    ea.BaseSystemCollection, ea.Constraints, ea.Forcing, ea.Contact
 ):
     pass
 
@@ -78,12 +78,15 @@ def simulate_rolling_friction_torque_with(C_s=0.0):
     static_mu_array = np.array([0.4, 0.4, 0.4])  # [forward, backward, sideways]
     kinetic_mu_array = np.array([0.2, 0.2, 0.2])  # [forward, backward, sideways]
 
-    rolling_friction_torque_sim.add_forcing_to(shearable_rod).using(
-        ea.AnisotropicFrictionalPlane,
+    friction_plane = ea.Plane(plane_origin=origin_plane, plane_normal=normal_plane)
+    rolling_friction_torque_sim.append(friction_plane)
+
+    rolling_friction_torque_sim.detect_contact_between(
+        shearable_rod, friction_plane
+    ).using(
+        ea.RodPlaneContactWithAnisotropicFriction,
         k=10.0,
         nu=1e-4,
-        plane_origin=origin_plane,
-        plane_normal=normal_plane,
         slip_velocity_tol=slip_velocity_tol,
         static_mu_array=static_mu_array,
         kinetic_mu_array=kinetic_mu_array,
diff --git a/examples/MuscularSnake/muscular_snake.py b/examples/MuscularSnake/muscular_snake.py
index 9aef0925..4741b37b 100644
--- a/examples/MuscularSnake/muscular_snake.py
+++ b/examples/MuscularSnake/muscular_snake.py
@@ -20,6 +20,7 @@ class MuscularSnakeSimulator(
     ea.Forcing,
     ea.CallBacks,
     ea.Damping,
+    ea.Contact,
 ):
     pass
 
@@ -346,12 +347,13 @@ class MuscularSnakeSimulator(
     [1.0 * mu, 1.5 * mu, 2.0 * mu]
 )  # [forward, backward, sideways]
 static_mu_array = 2 * kinetic_mu_array
-muscular_snake_simulator.add_forcing_to(snake_body).using(
-    ea.AnisotropicFrictionalPlane,
+friction_plane = ea.Plane(plane_origin=origin_plane, plane_normal=normal_plane)
+muscular_snake_simulator.append(friction_plane)
+
+muscular_snake_simulator.detect_contact_between(snake_body, friction_plane).using(
+    ea.RodPlaneContactWithAnisotropicFriction,
     k=1e1,
     nu=40,
-    plane_origin=origin_plane,
-    plane_normal=normal_plane,
     slip_velocity_tol=slip_velocity_tol,
     static_mu_array=static_mu_array,
     kinetic_mu_array=kinetic_mu_array,

From 71716c2585d67be4ebe558851502ac8001d519e4 Mon Sep 17 00:00:00 2001
From: Ali-7800 <47090295+Ali-7800@users.noreply.github.com>
Date: Tue, 14 Nov 2023 20:02:28 -0600
Subject: [PATCH 2/9] New example: ContinuumSnakeWithLiftingWaveCase

---
 .../continuum_snake_postprocessing.py         | 187 ++++++++++
 .../continuum_snake_with_lifting_wave.py      | 314 ++++++++++++++++
 .../snake_contact.py                          | 346 ++++++++++++++++++
 .../snake_forcing.py                          | 218 +++++++++++
 4 files changed, 1065 insertions(+)
 create mode 100644 examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_postprocessing.py
 create mode 100755 examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_with_lifting_wave.py
 create mode 100755 examples/ContinuumSnakeWithLiftingWaveCase/snake_contact.py
 create mode 100755 examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py

diff --git a/examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_postprocessing.py b/examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_postprocessing.py
new file mode 100644
index 00000000..37ad5ab9
--- /dev/null
+++ b/examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_postprocessing.py
@@ -0,0 +1,187 @@
+import numpy as np
+from matplotlib import pyplot as plt
+from matplotlib.colors import to_rgb
+from tqdm import tqdm
+
+
+def plot_snake_velocity(
+    plot_params: dict,
+    period,
+    filename="slithering_snake_velocity.png",
+    SAVE_FIGURE=False,
+):
+    time_per_period = np.array(plot_params["time"]) / period
+    avg_velocity = np.array(plot_params["avg_velocity"])
+
+    [
+        velocity_in_direction_of_rod,
+        velocity_in_rod_roll_dir,
+        _,
+        _,
+    ] = compute_projected_velocity(plot_params, period)
+
+    fig = plt.figure(figsize=(10, 8), frameon=True, dpi=150)
+    plt.rcParams.update({"font.size": 16})
+    ax = fig.add_subplot(111)
+    ax.grid(visible=True, which="minor", color="k", linestyle="--")
+    ax.grid(visible=True, which="major", color="k", linestyle="-")
+    ax.plot(
+        time_per_period[:], velocity_in_direction_of_rod[:, 2], "r-", label="forward"
+    )
+    ax.plot(
+        time_per_period[:],
+        velocity_in_rod_roll_dir[:, 0],
+        c=to_rgb("xkcd:bluish"),
+        label="lateral",
+    )
+    ax.plot(time_per_period[:], avg_velocity[:, 1], "k-", label="normal")
+    ax.set_ylabel("Velocity [m/s]", fontsize=16)
+    ax.set_xlabel("Time [s]", fontsize=16)
+    fig.legend(prop={"size": 20})
+    plt.show()
+    # Be a good boy and close figures
+    # https://stackoverflow.com/a/37451036
+    # plt.close(fig) alone does not suffice
+    # See https://github.com/matplotlib/matplotlib/issues/8560/
+    plt.close(plt.gcf())
+
+    if SAVE_FIGURE:
+        fig.savefig(filename)
+
+
+def plot_video(
+    plot_params: dict,
+    video_name="video.mp4",
+    fps=15,
+    xlim=(0, 4),
+    ylim=(-1, 1),
+):  # (time step, x/y/z, node)
+    import matplotlib.animation as manimation
+
+    positions_over_time = np.array(plot_params["position"])
+
+    print("plot video")
+    FFMpegWriter = manimation.writers["ffmpeg"]
+    metadata = dict(title="Movie Test", artist="Matplotlib", comment="Movie support!")
+    writer = FFMpegWriter(fps=fps, metadata=metadata)
+    fig = plt.figure(figsize=(10, 8), frameon=True, dpi=150)
+    ax = fig.add_subplot(111)
+    ax.set_xlim(*xlim)
+    ax.set_ylim(*ylim)
+    ax.set_xlabel("z [m]", fontsize=16)
+    ax.set_ylabel("x [m]", fontsize=16)
+    rod_lines_2d = ax.plot(positions_over_time[0][0], positions_over_time[0][1])[0]
+    # plt.axis("equal")
+    with writer.saving(fig, video_name, dpi=150):
+        for time in tqdm(range(1, len(plot_params["time"]))):
+            rod_lines_2d.set_xdata(positions_over_time[time][0])
+            rod_lines_2d.set_ydata(positions_over_time[time][1])
+            writer.grab_frame()
+
+    # Be a good boy and close figures
+    # https://stackoverflow.com/a/37451036
+    # plt.close(fig) alone does not suffice
+    # See https://github.com/matplotlib/matplotlib/issues/8560/
+    plt.close(plt.gcf())
+
+
+def compute_projected_velocity(plot_params: dict, period):
+
+    time_per_period = np.array(plot_params["time"]) / period
+    avg_velocity = np.array(plot_params["avg_velocity"])
+    center_of_mass = np.array(plot_params["center_of_mass"])
+
+    # Compute rod velocity in rod direction. We need to compute that because,
+    # after snake starts to move it chooses an arbitrary direction, which does not
+    # have to be initial tangent direction of the rod. Thus we need to project the
+    # snake velocity with respect to its new tangent and roll direction, after that
+    # we will get the correct forward and lateral speed. After this projection
+    # lateral velocity of the snake has to be oscillating between + and - values with
+    # zero mean.
+
+    # Number of steps in one period.
+    period_step = int(1.0 / (time_per_period[-1] - time_per_period[-2]))
+    number_of_period = int(time_per_period[-1])
+
+    # Center of mass position averaged in one period
+    center_of_mass_averaged_over_one_period = np.zeros((number_of_period - 2, 3))
+    for i in range(1, number_of_period - 1):
+        # position of center of mass averaged over one period
+        center_of_mass_averaged_over_one_period[i - 1] = np.mean(
+            center_of_mass[(i + 1) * period_step : (i + 2) * period_step]
+            - center_of_mass[(i + 0) * period_step : (i + 1) * period_step],
+            axis=0,
+        )
+    # Average the rod directions over multiple periods and get the direction of the rod.
+    direction_of_rod = np.mean(center_of_mass_averaged_over_one_period, axis=0)
+    direction_of_rod /= np.linalg.norm(direction_of_rod, ord=2)
+
+    # Compute the projected rod velocity in the direction of the rod
+    velocity_mag_in_direction_of_rod = np.einsum(
+        "ji,i->j", avg_velocity, direction_of_rod
+    )
+    velocity_in_direction_of_rod = np.einsum(
+        "j,i->ji", velocity_mag_in_direction_of_rod, direction_of_rod
+    )
+
+    # Get the lateral or roll velocity of the rod after subtracting its projected
+    # velocity in the direction of rod
+    velocity_in_rod_roll_dir = avg_velocity - velocity_in_direction_of_rod
+
+    # Compute the average velocity over the simulation, this can be used for optimizing snake
+    # for fastest forward velocity. We start after first period, because of the ramping up happens
+    # in first period.
+    average_velocity_over_simulation = np.mean(
+        velocity_in_direction_of_rod[period_step * 2 :], axis=0
+    )
+
+    return (
+        velocity_in_direction_of_rod,
+        velocity_in_rod_roll_dir,
+        average_velocity_over_simulation[2],
+        average_velocity_over_simulation[0],
+    )
+
+
+def plot_curvature(
+    plot_params: dict,
+    rest_lengths,
+    period,
+    save_fig=False,
+    filename="continuum_snake_curvature",
+):
+    s = np.cumsum(rest_lengths)
+    L0 = s[-1]
+    s = s / L0
+    s = s[:-1].copy()
+    x = np.linspace(0, 1, 100)
+    curvature = np.array(plot_params["curvature"])
+    time = np.array(plot_params["time"])
+    peak_time = period * 0.125
+    dt = time[1] - time[0]
+    peak_idx = int(peak_time / (dt))
+    plt.rcParams.update({"font.size": 16})
+    fig = plt.figure(figsize=(10, 8), frameon=True, dpi=150)
+    ax = fig.add_subplot(111)
+    try:
+        for i in range(peak_idx * 8, peak_idx * 8 * 2, peak_idx):
+            ax.plot(s, curvature[i, 0, :] * L0, "k")
+    except:
+        print("Simulation time not long enough to plot curvature")
+    ax.plot(
+        x, 7 * np.cos(2 * np.pi * x - 0.80), "--", label="stereotypical snake curvature"
+    )
+    ax.set_ylabel(r"$\kappa$", fontsize=16)
+    ax.set_xlabel("s", fontsize=16)
+    ax.set_xlim(0, 1)
+    ax.set_ylim(-10, 10)
+    fig.legend(prop={"size": 16})
+    plt.show()
+    if save_fig:
+        fig.savefig(filename)
+
+    # Be a good boy and close figures
+    # https://stackoverflow.com/a/37451036
+    # plt.close(fig) alone does not suffice
+    # See https://github.com/matplotlib/matplotlib/issues/8560/
+    plt.close(plt.gcf())
diff --git a/examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_with_lifting_wave.py b/examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_with_lifting_wave.py
new file mode 100755
index 00000000..96ae0fae
--- /dev/null
+++ b/examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_with_lifting_wave.py
@@ -0,0 +1,314 @@
+__doc__ = """Snake friction case from X. Zhang et. al. Nat. Comm. 2021"""
+
+import os
+import numpy as np
+import argparse
+import pickle
+
+from elastica import *
+
+from continuum_snake_postprocessing import (
+    plot_snake_velocity,
+    plot_video,
+    compute_projected_velocity,
+    plot_curvature,
+)
+from snake_forcing import (
+    MuscleTorques_snake,
+)
+from snake_contact import SnakeRodPlaneContact
+
+
+class SnakeSimulator(
+    BaseSystemCollection, Constraints, Forcing, Damping, CallBacks, Contact
+):
+    pass
+
+
+def run_snake(
+    b_coeff_lat,
+    PLOT_FIGURE=False,
+    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
+    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]
+
+    start = np.array([0.0, 0.0, 0.0 + base_radius])
+    direction = np.array([1.0, 0.0, 0.0])
+    normal = np.array([0.0, 0.0, 1.0])
+    density = 1000
+    E = 1e6
+    poisson_ratio = 0.5
+    shear_modulus = E / (poisson_ratio + 1.0)
+
+    shearable_rod = CosseratRod.straight_rod(
+        n_elem,
+        start,
+        direction,
+        normal,
+        base_length,
+        base_radius,
+        density,
+        youngs_modulus=E,
+        shear_modulus=shear_modulus,
+    )
+
+    snake_sim.append(shearable_rod)
+    damping_constant = phase_space_params.damping
+
+    # use linear damping with constant damping ratio
+    snake_sim.dampen(shearable_rod).using(
+        AnalyticalLinearDamper,
+        damping_constant=damping_constant,
+        time_step=time_step,
+    )
+
+    # Add gravitational forces
+    gravitational_acc = -9.80665
+
+    snake_sim.add_forcing_to(shearable_rod).using(
+        GravityForces, acc_gravity=np.array([0.0, 0.0, gravitational_acc])
+    )
+
+    # 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_amp = b_coeff_lat[:-1]
+
+    lateral_ratio = 1.0  # switch of lateral wave
+    snake_sim.add_forcing_to(shearable_rod).using(
+        MuscleTorques_snake,
+        base_length=base_length,
+        b_coeff=snake_torque_ratio * lateral_ratio * lateral_amp,
+        period=period,
+        wave_number=2.0 * np.pi / (lateral_wave_length),
+        phase_shift=0.0,
+        rest_lengths=shearable_rod.rest_lengths,
+        ramp_up_time=period,
+        direction=normal,
+        with_spline=True,
+    )
+
+    # 2. Add muscle torques -- lifting wave
+    # Define lifting wave parameters
+    lift_wave_length = lateral_wave_length
+    lift_amp = np.array([1e-3, 2e-3, 2e-3, 2e-3, 2e-3, 1e-3])
+
+    lift_ratio = 1.0  # switch of lifting wave
+    phase = 0.5
+    snake_sim.add_forcing_to(shearable_rod).using(
+        MuscleTorques_snake,
+        base_length=base_length,
+        b_coeff=snake_torque_liftratio * lift_ratio * lift_amp,
+        period=period,
+        wave_number=2.0 * np.pi / (lift_wave_length),
+        phase_shift=phase * period,
+        rest_lengths=shearable_rod.rest_lengths,
+        ramp_up_time=0.01,
+        direction=normal,
+        with_spline=True,
+        switch_on_time=2.0,
+        is_lateral_wave=False,
+    )
+
+    # Some common parameters first - define friction ratio etc.
+    slip_velocity_tol = 1e-8
+    froude = 0.1
+    mu = base_length / (period * period * np.abs(gravitational_acc) * froude)
+    kinetic_mu_array = np.array(
+        [mu, 1.5 * mu, 2.0 * mu]
+    )  # [forward, backward, sideways]
+    normal_plane = normal
+    origin_plane = np.array([0.0, 0.0, 0.0])
+    ground_plane = Plane(plane_normal=normal_plane, plane_origin=origin_plane)
+    snake_sim.append(ground_plane)
+
+    snake_sim.detect_contact_between(shearable_rod, ground_plane).using(
+        SnakeRodPlaneContact,
+        k=1e2,
+        nu=1e-1,
+        slip_velocity_tol=slip_velocity_tol,
+        kinetic_mu_array=kinetic_mu_array,
+    )
+
+    # Add call backs
+    class ContinuumSnakeCallBack(CallBackBaseClass):
+        """
+        Call back function for continuum snake
+        """
+
+        def __init__(self, step_skip: int, callback_params: dict):
+            CallBackBaseClass.__init__(self)
+            self.every = step_skip
+            self.callback_params = callback_params
+
+        def make_callback(self, system, time, current_step: int):
+
+            if current_step % self.every == 0:
+
+                self.callback_params["time"].append(time)
+                self.callback_params["step"].append(current_step)
+                self.callback_params["position"].append(
+                    system.position_collection.copy()
+                )
+                self.callback_params["radius"].append(system.radius.copy())
+                self.callback_params["velocity"].append(
+                    system.velocity_collection.copy()
+                )
+                self.callback_params["avg_velocity"].append(
+                    system.compute_velocity_center_of_mass()
+                )
+
+                self.callback_params["center_of_mass"].append(
+                    system.compute_position_center_of_mass()
+                )
+                self.callback_params["curvature"].append(system.kappa.copy())
+
+                return
+
+    pp_list = defaultdict(list)
+    snake_sim.collect_diagnostics(shearable_rod).using(
+        ContinuumSnakeCallBack, step_skip=step_skip, callback_params=pp_list
+    )
+
+    snake_sim.finalize()
+
+    timestepper = PositionVerlet()
+    integrate(timestepper, snake_sim, final_time, total_steps)
+
+    if PLOT_FIGURE:
+        filename_plot = "continuum_snake_velocity.png"
+        plot_snake_velocity(pp_list, period, filename_plot, SAVE_FIGURE)
+        plot_curvature(pp_list, shearable_rod.rest_lengths, period, SAVE_FIGURE)
+
+        if SAVE_VIDEO:
+            filename_video = "continuum_snake_with_lifting_wave.mp4"
+            plot_video(
+                pp_list,
+                video_name=filename_video,
+                fps=rendering_fps,
+                xlim=(0, 3),
+                ylim=(-1, 1),
+            )
+
+    if SAVE_RESULTS:
+
+        filename = "continuum_snake.dat"
+        file = open(filename, "wb")
+        pickle.dump(pp_list, file)
+        file.close()
+
+    # Compute the average forward velocity. These will be used for optimization.
+    [_, _, avg_forward, avg_lateral] = compute_projected_velocity(pp_list, period)
+
+    return avg_forward, avg_lateral, pp_list
+
+
+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
+    SAVE_VIDEO = True
+    SAVE_RESULTS = True
+    CMA_OPTION = False
+
+    if CMA_OPTION:
+        import cma
+
+        SAVE_OPTIMIZED_COEFFICIENTS = False
+
+        def optimize_snake(spline_coefficient):
+            [avg_forward, _, _] = run_snake(
+                spline_coefficient,
+                PLOT_FIGURE=False,
+                SAVE_FIGURE=False,
+                SAVE_VIDEO=False,
+                SAVE_RESULTS=False,
+            )
+            return -avg_forward
+
+        # Optimize snake for forward velocity. In cma.fmin first input is function
+        # to be optimized, second input is initial guess for coefficients you are optimizing
+        # for and third input is standard deviation you initially set.
+        optimized_spline_coefficients = cma.fmin(optimize_snake, 7 * [0], 0.5)
+
+        # Save the optimized coefficients to a file
+        filename_data = "optimized_coefficients.txt"
+        if SAVE_OPTIMIZED_COEFFICIENTS:
+            assert filename_data != "", "provide a file name for coefficients"
+            np.savetxt(filename_data, optimized_spline_coefficients, delimiter=",")
+
+    else:
+        # Add muscle forces on the rod
+        if os.path.exists("optimized_coefficients.txt"):
+            t_coeff_optimized = np.genfromtxt(
+                "optimized_coefficients.txt", delimiter=","
+            )
+        else:
+            wave_length = 1.0
+            t_coeff_optimized = np.array(
+                [4e-3, 4e-3, 4e-3, 4e-3, 4e-3, 4e-3]
+                # [3.4e-3, 3.3e-3, 5.7e-3, 2.8e-3, 3.0e-3, 3.0e-3]
+            )
+            t_coeff_optimized = np.hstack((t_coeff_optimized, wave_length))
+
+        # run the simulation
+        [avg_forward, avg_lateral, pp_list] = run_snake(
+            t_coeff_optimized, PLOT_FIGURE, SAVE_FIGURE, SAVE_VIDEO, SAVE_RESULTS, args
+        )
+
+        print("average forward velocity:", avg_forward)
+        print("average forward lateral:", avg_lateral)
diff --git a/examples/ContinuumSnakeWithLiftingWaveCase/snake_contact.py b/examples/ContinuumSnakeWithLiftingWaveCase/snake_contact.py
new file mode 100755
index 00000000..e2c0da83
--- /dev/null
+++ b/examples/ContinuumSnakeWithLiftingWaveCase/snake_contact.py
@@ -0,0 +1,346 @@
+__doc__ = """Rod plane contact with anistropic friction (no static friction)"""
+
+import numpy as np
+from elastica._linalg import (
+    _batch_norm,
+    _batch_product_i_k_to_ik,
+    _batch_product_k_ik_to_ik,
+    _batch_vec_oneD_vec_cross,
+    _batch_matvec,
+    _batch_product_i_ik_to_k,
+    _batch_dot,
+    _batch_matrix_transpose,
+    _batch_cross,
+    _batch_vector_sum,
+)
+
+from elastica.contact_utils import (
+    _node_to_element_position,
+    _node_to_element_velocity,
+    _find_slipping_elements,
+    _elements_to_nodes_inplace,
+    _node_to_element_mass_or_force,
+)
+from numba import njit
+from elastica.rod import RodBase
+from elastica.surface import Plane
+from elastica.contact_forces import NoContact
+from elastica.typing import RodType, SystemType, AllowedContactType
+
+
+@njit(cache=True)
+def apply_normal_force_numba(
+    plane_origin,
+    plane_normal,
+    surface_tol,
+    k,
+    nu,
+    radius,
+    mass,
+    position_collection,
+    velocity_collection,
+    internal_forces,
+    external_forces,
+):
+    """
+    This function computes the plane force response on the element, in the
+    case of contact. Contact model given in Eqn 4.8 Gazzola et. al. RSoS 2018 paper
+    is used.
+
+    Parameters
+    ----------
+    system
+
+    Returns
+    -------
+    magnitude of the plane response
+    """
+
+    # Compute plane response force
+    nodal_total_forces = _batch_vector_sum(internal_forces, external_forces)
+    element_total_forces = _node_to_element_mass_or_force(nodal_total_forces)
+
+    force_component_along_normal_direction = _batch_product_i_ik_to_k(
+        plane_normal, element_total_forces
+    )
+    forces_along_normal_direction = _batch_product_i_k_to_ik(
+        plane_normal, force_component_along_normal_direction
+    )
+
+    # If the total force component along the plane normal direction is greater than zero that means,
+    # total force is pushing rod away from the plane not towards the plane. Thus, response force
+    # applied by the surface has to be zero.
+    forces_along_normal_direction[
+        ..., np.where(force_component_along_normal_direction > 0)[0]
+    ] = 0.0
+    # Compute response force on the element. Plane response force
+    # has to be away from the surface and towards the element. Thus
+    # multiply forces along normal direction with negative sign.
+    plane_response_force = -forces_along_normal_direction
+
+    # Elastic force response due to penetration
+    element_position = _node_to_element_position(position_collection)
+    distance_from_plane = _batch_product_i_ik_to_k(
+        plane_normal, (element_position - plane_origin)
+    )
+    plane_penetration = np.minimum(distance_from_plane - radius, 0.0)
+    elastic_force = -k * _batch_product_i_k_to_ik(plane_normal, plane_penetration)
+
+    # Damping force response due to velocity towards the plane
+    element_velocity = _node_to_element_velocity(
+        mass=mass, node_velocity_collection=velocity_collection
+    )
+    normal_component_of_element_velocity = _batch_product_i_ik_to_k(
+        plane_normal, element_velocity
+    )
+    damping_force = -nu * _batch_product_i_k_to_ik(
+        plane_normal, normal_component_of_element_velocity
+    )
+
+    # Compute total plane response force
+    plane_response_force_total = elastic_force + damping_force
+
+    # Check if the rod elements are in contact with plane.
+    no_contact_point_idx = np.where((distance_from_plane - radius) > surface_tol)[0]
+    # If rod element does not have any contact with plane, plane cannot apply response
+    # force on the element. Thus lets set plane response force to 0.0 for the no contact points.
+    plane_response_force[..., no_contact_point_idx] = 0.0
+    plane_response_force_total[..., no_contact_point_idx] = 0.0
+
+    # Update the external forces
+    _elements_to_nodes_inplace(plane_response_force_total, external_forces)
+
+    return (_batch_norm(plane_response_force), no_contact_point_idx)
+
+
+@njit(cache=True)
+def anisotropic_friction(
+    plane_origin,
+    plane_normal,
+    surface_tol,
+    slip_velocity_tol,
+    k,
+    nu,
+    kinetic_mu_forward,
+    kinetic_mu_backward,
+    kinetic_mu_sideways,
+    radius,
+    mass,
+    tangents,
+    position_collection,
+    director_collection,
+    velocity_collection,
+    omega_collection,
+    internal_forces,
+    external_forces,
+):
+    plane_response_force_mag, no_contact_point_idx = apply_normal_force_numba(
+        plane_origin,
+        plane_normal,
+        surface_tol,
+        k,
+        nu,
+        radius,
+        mass,
+        position_collection,
+        velocity_collection,
+        internal_forces,
+        external_forces,
+    )
+
+    # First compute component of rod tangent in plane. Because friction forces acts in plane not out of plane. Thus
+    # axial direction has to be in plane, it cannot be out of plane. We are projecting rod element tangent vector in
+    # to the plane. So friction forces can only be in plane forces and not out of plane.
+    tangent_along_normal_direction = _batch_product_i_ik_to_k(plane_normal, tangents)
+    tangent_perpendicular_to_normal_direction = tangents - _batch_product_i_k_to_ik(
+        plane_normal, tangent_along_normal_direction
+    )
+    tangent_perpendicular_to_normal_direction_mag = _batch_norm(
+        tangent_perpendicular_to_normal_direction
+    )
+    # Normalize tangent_perpendicular_to_normal_direction. This is axial direction for plane. Here we are adding
+    # small tolerance (1e-10) for normalization, in order to prevent division by 0.
+    axial_direction = _batch_product_k_ik_to_ik(
+        1 / (tangent_perpendicular_to_normal_direction_mag + 1e-14),
+        tangent_perpendicular_to_normal_direction,
+    )
+    element_velocity = _node_to_element_velocity(
+        mass=mass, node_velocity_collection=velocity_collection
+    )
+    # first apply axial kinetic friction
+    velocity_mag_along_axial_direction = _batch_dot(element_velocity, axial_direction)
+    velocity_along_axial_direction = _batch_product_k_ik_to_ik(
+        velocity_mag_along_axial_direction, axial_direction
+    )
+
+    # Friction forces depends on the direction of velocity, in other words sign
+    # of the velocity vector.
+    velocity_sign_along_axial_direction = np.sign(velocity_mag_along_axial_direction)
+    # Check top for sign convention
+    kinetic_mu = 0.5 * (
+        kinetic_mu_forward * (1 + velocity_sign_along_axial_direction)
+        + kinetic_mu_backward * (1 - velocity_sign_along_axial_direction)
+    )
+    # Call slip function to check if elements slipping or not
+    slip_function_along_axial_direction = _find_slipping_elements(
+        velocity_along_axial_direction, slip_velocity_tol
+    )
+
+    # Now rolling kinetic friction
+    rolling_direction = _batch_vec_oneD_vec_cross(axial_direction, plane_normal)
+    torque_arm = _batch_product_i_k_to_ik(-plane_normal, radius)
+    velocity_along_rolling_direction = _batch_dot(element_velocity, rolling_direction)
+    directors_transpose = _batch_matrix_transpose(director_collection)
+    # w_rot = Q.T @ omega @ Q @ r
+    rotation_velocity = _batch_matvec(
+        directors_transpose,
+        _batch_cross(omega_collection, _batch_matvec(director_collection, torque_arm)),
+    )
+    rotation_velocity_along_rolling_direction = _batch_dot(
+        rotation_velocity, rolling_direction
+    )
+    slip_velocity_mag_along_rolling_direction = (
+        velocity_along_rolling_direction + rotation_velocity_along_rolling_direction
+    )
+    slip_velocity_along_rolling_direction = _batch_product_k_ik_to_ik(
+        slip_velocity_mag_along_rolling_direction, rolling_direction
+    )
+    slip_function_along_rolling_direction = _find_slipping_elements(
+        slip_velocity_along_rolling_direction, slip_velocity_tol
+    )
+
+    unitized_total_velocity = element_velocity / _batch_norm(element_velocity + 1e-14)
+    # Apply kinetic friction in axial direction.
+    kinetic_friction_force_along_axial_direction = -(
+        (1.0 - slip_function_along_axial_direction)
+        * kinetic_mu
+        * plane_response_force_mag
+        * _batch_dot(unitized_total_velocity, axial_direction)
+        * axial_direction
+    )
+    # If rod element does not have any contact with plane, plane cannot apply friction
+    # force on the element. Thus lets set kinetic friction force to 0.0 for the no contact points.
+    kinetic_friction_force_along_axial_direction[..., no_contact_point_idx] = 0.0
+    _elements_to_nodes_inplace(
+        kinetic_friction_force_along_axial_direction, external_forces
+    )
+    # Apply kinetic friction in rolling direction.
+    kinetic_friction_force_along_rolling_direction = -(
+        (1.0 - slip_function_along_rolling_direction)
+        * kinetic_mu_sideways
+        * plane_response_force_mag
+        * _batch_dot(unitized_total_velocity, rolling_direction)
+        * rolling_direction
+    )
+    # If rod element does not have any contact with plane, plane cannot apply friction
+    # force on the element. Thus lets set kinetic friction force to 0.0 for the no contact points.
+    kinetic_friction_force_along_rolling_direction[..., no_contact_point_idx] = 0.0
+    _elements_to_nodes_inplace(
+        kinetic_friction_force_along_rolling_direction, external_forces
+    )
+
+
+class SnakeRodPlaneContact(NoContact):
+    """
+    This class is for applying contact forces between a snake rod and a plane with friction.
+    First system is always rod and second system is always plane.
+
+    How to define contact between rod and plane.
+    >>> simulator.detect_contact_between(rod, plane).using(
+    ...    SnakeRodPlaneContact,
+    ...    k=1e4,
+    ...    nu=10,
+    ...    slip_velocity_tol = 1e-4,
+    ...    kinetic_mu_array = np.array([1.0,2.0,3.0]),
+    ... )
+    """
+
+    def __init__(
+        self,
+        k: float,
+        nu: float,
+        slip_velocity_tol: float,
+        kinetic_mu_array: np.ndarray,
+    ):
+        """
+        Parameters
+        ----------
+        k : float
+                Contact spring constant.
+        nu : float
+                Contact damping constant.
+        slip_velocity_tol: float
+                Velocity tolerance to determine if the element is slipping or not.
+        kinetic_mu_array: numpy.ndarray
+                1D (3,) array containing data with 'float' type.
+                [forward, backward, sideways] kinetic friction coefficients.
+        """
+        super(SnakeRodPlaneContact, self).__init__()
+        self.k = k
+        self.nu = nu
+        self.surface_tol = 1e-4
+        self.slip_velocity_tol = slip_velocity_tol
+        (
+            self.kinetic_mu_forward,
+            self.kinetic_mu_backward,
+            self.kinetic_mu_sideways,
+        ) = kinetic_mu_array
+
+    def _check_systems_validity(
+        self,
+        system_one: SystemType,
+        system_two: AllowedContactType,
+    ) -> None:
+        """
+        This checks the contact order and type of a SystemType object and an AllowedContactType object.
+        For the RodSphereContact class first_system should be a rod and second_system should be a plane.
+        Parameters
+        ----------
+        system_one
+                SystemType
+        system_two
+                AllowedContactType
+        """
+        if not issubclass(system_one.__class__, RodBase) or not issubclass(
+            system_two.__class__, Plane
+        ):
+            raise TypeError(
+                "Systems provided to the contact class have incorrect order/type. \n"
+                " First system is {0} and second system is {1}. \n"
+                " First system should be a rod, second should be a plane".format(
+                    system_one.__class__, system_two.__class__
+                )
+            )
+
+    def apply_contact(self, system_one: RodType, system_two: SystemType) -> None:
+        """
+        In the case of contact with the plane, this function computes the plane reaction force on the element.
+
+        Parameters
+        ----------
+        system_one: object
+                Rod object.
+        system_two: object
+                Plane object.
+
+        """
+        anisotropic_friction(
+            system_two.origin,
+            system_two.normal,
+            self.surface_tol,
+            self.slip_velocity_tol,
+            self.k,
+            self.nu,
+            self.kinetic_mu_forward,
+            self.kinetic_mu_backward,
+            self.kinetic_mu_sideways,
+            system_one.radius,
+            system_one.mass,
+            system_one.tangents,
+            system_one.position_collection,
+            system_one.director_collection,
+            system_one.velocity_collection,
+            system_one.omega_collection,
+            system_one.internal_forces,
+            system_one.external_forces,
+        )
diff --git a/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py b/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py
new file mode 100755
index 00000000..e6d1304c
--- /dev/null
+++ b/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py
@@ -0,0 +1,218 @@
+__doc__ = """ External forces and Actions (in learning) of the snake terrain case."""
+
+import numpy as np
+from elastica._linalg import _batch_matvec
+from elastica.utils import _bspline
+from numba import njit
+from elastica._linalg import (
+    _batch_norm,
+    _batch_product_i_k_to_ik,
+    _batch_product_k_ik_to_ik,
+    _batch_vec_oneD_vec_cross,
+)
+from elastica.external_forces import NoForces
+from elastica.external_forces import (
+    inplace_addition,
+    inplace_substraction,
+)
+
+
+class MuscleTorques_snake(NoForces):
+    """
+    This class applies muscle torques along the body. The applied muscle torques are treated
+    as applied external forces. This class can apply
+    muscle torques as a traveling wave with a beta spline or only
+    as a traveling wave. For implementation details refer to Gazzola et. al.
+    RSoS. (2018).
+
+            Attributes
+            ----------
+            direction: numpy.ndarray
+                    2D (dim, 1) array containing data with 'float' type. Muscle torque direction.
+            angular_frequency: float
+                    Angular frequency of traveling wave.
+            wave_number: float
+                    Wave number of traveling wave.
+            phase_shift: float
+                    Phase shift of traveling wave.
+            ramp_up_time: float
+                    Applied muscle torques are ramped up until ramp up time.
+            my_spline: numpy.ndarray
+                    1D (blocksize) array containing data with 'float' type. Generated spline.
+
+            NEW PARAMETER:
+            switch_on_time: float
+                    time to switch on the muscle activation.
+            is_lateral_wave: bool
+                    check if it is lateral muscle torque.
+    """
+
+    def __init__(
+        self,
+        base_length,
+        b_coeff,
+        period,
+        wave_number,
+        phase_shift,
+        direction,
+        rest_lengths,
+        ramp_up_time=0.0,
+        with_spline=False,
+        switch_on_time=0.0,
+        is_lateral_wave=True,
+        *args,
+        **kwargs,
+    ):
+        """
+
+        Parameters
+        ----------
+        base_length: float
+                Rest length of the rod-like object.
+        b_coeff: nump.ndarray
+                1D array containing data with 'float' type.
+                Beta coefficients for beta-spline.
+        period: float
+                Period of traveling wave.
+        wave_number: float
+                Wave number of traveling wave.
+        phase_shift: float
+                Phase shift of traveling wave.
+        direction: numpy.ndarray
+           1D (dim) array containing data with 'float' type.
+        ramp_up_time: float
+                Applied muscle torques are ramped up until ramp up time.
+        with_spline: boolean
+                Option to use beta-spline.
+
+        NEW PARAMETER:
+        switch_on_time: float
+                time to switch on the muscle activation.
+        is_lateral_wave: bool
+                check if it is lateral muscle torque.
+
+        """
+        super(MuscleTorques_snake, self).__init__()
+
+        self.direction = direction  # Direction torque applied
+        self.angular_frequency = 2.0 * np.pi / period
+        self.wave_number = wave_number
+        self.phase_shift = phase_shift
+
+        assert ramp_up_time >= 0.0
+        self.ramp_up_time = ramp_up_time
+        assert switch_on_time >= 0.0
+        self.switch_on_time = switch_on_time
+        self.is_lateral_wave = is_lateral_wave
+
+        # s is the position of nodes on the rod, we go from node=1 to node=nelem-1, because there is no
+        # torques applied by first and last node on elements. Reason is that we cannot apply torque in an
+        # infinitesimal segment at the beginning and end of rod, because there is no additional element
+        # (at element=-1 or element=n_elem+1) to provide internal torques to cancel out an external
+        # torque. This coupled with the requirement that the sum of all muscle torques has
+        # to be zero results in this condition.
+        self.s = np.cumsum(rest_lengths)
+        self.s /= self.s[-1]
+
+        if with_spline:
+            assert b_coeff.size != 0, "Beta spline coefficient array (t_coeff) is empty"
+            my_spline, ctr_pts, ctr_coeffs = _bspline(b_coeff)
+            self.my_spline = my_spline(self.s)
+
+        else:
+
+            def constant_function(input):
+                """
+                Return array of ones same as the size of the input array. This
+                function is called when Beta spline function is not used.
+
+                Parameters
+                ----------
+                input
+
+                Returns
+                -------
+
+                """
+                return np.ones(input.shape)
+
+            self.my_spline = constant_function(self.s)
+
+    def apply_torques(self, system, time: np.float64 = 0.0):
+        self.compute_muscle_torques(
+            time,
+            self.my_spline,
+            self.s,
+            self.angular_frequency,
+            self.wave_number,
+            self.phase_shift,
+            self.ramp_up_time,
+            self.direction,
+            self.switch_on_time,
+            self.is_lateral_wave,
+            system.tangents,
+            system.director_collection,
+            system.external_torques,
+        )
+
+    @staticmethod
+    @njit(cache=True)
+    def compute_muscle_torques(
+        time,
+        my_spline,
+        s,
+        angular_frequency,
+        wave_number,
+        phase_shift,
+        ramp_up_time,
+        direction,
+        switch_on_time,
+        is_lateral_wave,
+        tangents,
+        director_collection,
+        external_torques,
+    ):
+        if time > switch_on_time:
+            # Ramp up the muscle torque
+            factor = min(1.0, (time - switch_on_time) / ramp_up_time)
+            # From the node 1 to node nelem-1
+            # Magnitude of the torque. Am = beta(s) * sin(2pi*t/T + 2pi*s/lambda + phi)
+            # There is an inconsistency with paper and Elastica cpp implementation. In paper sign in
+            # front of wave number is positive, in Elastica cpp it is negative.
+            torque_mag = (
+                factor
+                * my_spline
+                * np.sin(
+                    angular_frequency * (time - switch_on_time - phase_shift)
+                    - wave_number * s
+                )
+            )
+            # 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
+                )
+
+            inplace_addition(
+                external_torques[..., 1:],
+                _batch_matvec(director_collection[..., 1:], torque),
+            )
+            inplace_substraction(
+                external_torques[..., :-1],
+                _batch_matvec(director_collection[..., :-1], torque),
+            )

From 944b7351478180e676530451e2deb2a6b4ff4f02 Mon Sep 17 00:00:00 2001
From: Ali-7800 <47090295+Ali-7800@users.noreply.github.com>
Date: Tue, 28 Nov 2023 01:13:46 -0600
Subject: [PATCH 3/9] fixed typo

---
 elastica/contact_forces.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/elastica/contact_forces.py b/elastica/contact_forces.py
index 83657ce6..8f3fdc6e 100644
--- a/elastica/contact_forces.py
+++ b/elastica/contact_forces.py
@@ -539,7 +539,7 @@ def _check_systems_validity(
     ) -> None:
         """
         This checks the contact order and type of a SystemType object and an AllowedContactType object.
-        For the RodSphereContact class first_system should be a rod and second_system should be a plane.
+        For the RodPlaneContact class first_system should be a rod and second_system should be a plane.
         Parameters
         ----------
         system_one

From a15d811d804082951bc3293f08038398b87a498f Mon Sep 17 00:00:00 2001
From: Ali-7800 <47090295+Ali-7800@users.noreply.github.com>
Date: Tue, 28 Nov 2023 01:15:13 -0600
Subject: [PATCH 4/9] removed snake_id options and using native MuscularTorques
 for lateral wave

---
 .../continuum_snake_with_lifting_wave.py      | 60 ++++---------------
 .../snake_forcing.py                          | 37 ++++++------
 2 files changed, 29 insertions(+), 68 deletions(-)

diff --git a/examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_with_lifting_wave.py b/examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_with_lifting_wave.py
index 96ae0fae..52b37586 100755
--- a/examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_with_lifting_wave.py
+++ b/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)
diff --git a/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py b/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py
index e6d1304c..3e13d8fb 100755
--- a/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py
+++ b/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:],

From e10cc4b9537951092bf6d17e542aa16d311f73c3 Mon Sep 17 00:00:00 2001
From: Ali-7800 <47090295+Ali-7800@users.noreply.github.com>
Date: Tue, 28 Nov 2023 01:23:34 -0600
Subject: [PATCH 5/9] removed unused import

---
 examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py b/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py
index 3e13d8fb..8d23f361 100755
--- a/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py
+++ b/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py
@@ -6,7 +6,6 @@
 from numba import njit
 from elastica._linalg import (
     _batch_norm,
-    _batch_product_i_k_to_ik,
     _batch_product_k_ik_to_ik,
     _batch_vec_oneD_vec_cross,
 )

From b01663eefdca5dd3aba396cffcf5e6c29ba4deb6 Mon Sep 17 00:00:00 2001
From: Ali-7800 <47090295+Ali-7800@users.noreply.github.com>
Date: Tue, 28 Nov 2023 01:31:51 -0600
Subject: [PATCH 6/9] removed unused imports

---
 .../continuum_snake_with_lifting_wave.py                         | 1 -
 1 file changed, 1 deletion(-)

diff --git a/examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_with_lifting_wave.py b/examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_with_lifting_wave.py
index 52b37586..1c275634 100755
--- a/examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_with_lifting_wave.py
+++ b/examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_with_lifting_wave.py
@@ -2,7 +2,6 @@
 
 import os
 import numpy as np
-import argparse
 import pickle
 
 from elastica import *

From 6d0c1af9257b7ef8264c1e47f35aa61ab49dc0e5 Mon Sep 17 00:00:00 2001
From: Ali-7800 <47090295+Ali-7800@users.noreply.github.com>
Date: Tue, 28 Nov 2023 16:25:30 -0600
Subject: [PATCH 7/9] updated docs and removed unused is_lateral_wave option

---
 .../continuum_snake_with_lifting_wave.py          |  1 -
 .../snake_forcing.py                              | 15 ++-------------
 2 files changed, 2 insertions(+), 14 deletions(-)

diff --git a/examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_with_lifting_wave.py b/examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_with_lifting_wave.py
index 1c275634..ba18b198 100755
--- a/examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_with_lifting_wave.py
+++ b/examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_with_lifting_wave.py
@@ -124,7 +124,6 @@ def run_snake(
         direction=normal,
         with_spline=True,
         switch_on_time=2.0,
-        is_lateral_wave=False,
     )
 
     # Some common parameters first - define friction ratio etc.
diff --git a/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py b/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py
index 8d23f361..3de1496a 100755
--- a/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py
+++ b/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py
@@ -19,10 +19,9 @@
 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
+    as applied external forces. This class can apply lifting
     muscle torques as a traveling wave with a beta spline or only
-    as a traveling wave. For implementation details refer to Gazzola et. al.
-    RSoS. (2018).
+    as a traveling wave. For implementation details refer to X. Zhang et. al. Nat. Comm. 2021
 
             Attributes
             ----------
@@ -38,12 +37,8 @@ class MuscleTorques_lifting(NoForces):
                     Applied muscle torques are ramped up until ramp up time.
             my_spline: numpy.ndarray
                     1D (blocksize) array containing data with 'float' type. Generated spline.
-
-            NEW PARAMETER:
             switch_on_time: float
                     time to switch on the muscle activation.
-            is_lateral_wave: bool
-                    check if it is lateral muscle torque.
     """
 
     def __init__(
@@ -58,7 +53,6 @@ def __init__(
         ramp_up_time=0.0,
         with_spline=False,
         switch_on_time=0.0,
-        is_lateral_wave=True,
         *args,
         **kwargs,
     ):
@@ -83,12 +77,8 @@ def __init__(
                 Applied muscle torques are ramped up until ramp up time.
         with_spline: boolean
                 Option to use beta-spline.
-
-        NEW PARAMETER:
         switch_on_time: float
                 time to switch on the muscle activation.
-        is_lateral_wave: bool
-                check if it is lateral muscle torque.
 
         """
         super(MuscleTorques_lifting, self).__init__()
@@ -102,7 +92,6 @@ def __init__(
         self.ramp_up_time = ramp_up_time
         assert switch_on_time >= 0.0
         self.switch_on_time = switch_on_time
-        self.is_lateral_wave = is_lateral_wave
 
         # s is the position of nodes on the rod, we go from node=1 to node=nelem-1, because there is no
         # torques applied by first and last node on elements. Reason is that we cannot apply torque in an

From c447ae1d4780cdd37e16be6f30ec031a6a6f13c7 Mon Sep 17 00:00:00 2001
From: Ali-7800 <47090295+Ali-7800@users.noreply.github.com>
Date: Tue, 28 Nov 2023 16:29:19 -0600
Subject: [PATCH 8/9] remove is_lateral_wave option

---
 examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py b/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py
index 3de1496a..31664e1c 100755
--- a/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py
+++ b/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py
@@ -137,7 +137,6 @@ def apply_torques(self, system, time: np.float64 = 0.0):
             self.ramp_up_time,
             self.direction,
             self.switch_on_time,
-            self.is_lateral_wave,
             system.tangents,
             system.director_collection,
             system.external_torques,
@@ -155,7 +154,6 @@ def compute_muscle_torques(
         ramp_up_time,
         direction,
         switch_on_time,
-        is_lateral_wave,
         tangents,
         director_collection,
         external_torques,

From 1273c13f42fde40dfd93858e7f06e666f6db39e4 Mon Sep 17 00:00:00 2001
From: Ali-7800 <47090295+Ali-7800@users.noreply.github.com>
Date: Sun, 3 Dec 2023 17:39:06 -0600
Subject: [PATCH 9/9] Minor name and style changes

---
 .../continuum_snake_with_lifting_wave.py      |  5 +--
 .../snake_contact.py                          |  2 +-
 .../snake_forcing.py                          | 37 ++++---------------
 3 files changed, 11 insertions(+), 33 deletions(-)

diff --git a/examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_with_lifting_wave.py b/examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_with_lifting_wave.py
index ba18b198..fec1fc90 100755
--- a/examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_with_lifting_wave.py
+++ b/examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_with_lifting_wave.py
@@ -13,7 +13,7 @@
     plot_curvature,
 )
 from snake_forcing import (
-    MuscleTorques_lifting,
+    MuscleTorquesLifting,
 )
 from snake_contact import SnakeRodPlaneContact
 
@@ -113,8 +113,7 @@ def run_snake(
     lift_ratio = 1.0  # switch of lifting wave
     phase = 0.5
     snake_sim.add_forcing_to(shearable_rod).using(
-        MuscleTorques_lifting,
-        base_length=base_length,
+        MuscleTorquesLifting,
         b_coeff=snake_torque_liftratio * lift_ratio * lift_amp,
         period=period,
         wave_number=2.0 * np.pi / (lift_wave_length),
diff --git a/examples/ContinuumSnakeWithLiftingWaveCase/snake_contact.py b/examples/ContinuumSnakeWithLiftingWaveCase/snake_contact.py
index e2c0da83..68dea81d 100755
--- a/examples/ContinuumSnakeWithLiftingWaveCase/snake_contact.py
+++ b/examples/ContinuumSnakeWithLiftingWaveCase/snake_contact.py
@@ -275,7 +275,7 @@ def __init__(
                 1D (3,) array containing data with 'float' type.
                 [forward, backward, sideways] kinetic friction coefficients.
         """
-        super(SnakeRodPlaneContact, self).__init__()
+        super().__init__()
         self.k = k
         self.nu = nu
         self.surface_tol = 1e-4
diff --git a/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py b/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py
index 31664e1c..3b471aca 100755
--- a/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py
+++ b/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py
@@ -16,7 +16,7 @@
 )
 
 
-class MuscleTorques_lifting(NoForces):
+class MuscleTorquesLifting(NoForces):
     """
     This class applies muscle torques along the body. The applied muscle torques are treated
     as applied external forces. This class can apply lifting
@@ -43,7 +43,6 @@ class MuscleTorques_lifting(NoForces):
 
     def __init__(
         self,
-        base_length,
         b_coeff,
         period,
         wave_number,
@@ -53,15 +52,11 @@ def __init__(
         ramp_up_time=0.0,
         with_spline=False,
         switch_on_time=0.0,
-        *args,
-        **kwargs,
     ):
         """
 
         Parameters
         ----------
-        base_length: float
-                Rest length of the rod-like object.
         b_coeff: nump.ndarray
                 1D array containing data with 'float' type.
                 Beta coefficients for beta-spline.
@@ -81,7 +76,7 @@ def __init__(
                 time to switch on the muscle activation.
 
         """
-        super(MuscleTorques_lifting, self).__init__()
+        super().__init__()
 
         self.direction = direction  # Direction torque applied
         self.angular_frequency = 2.0 * np.pi / period
@@ -104,32 +99,16 @@ def __init__(
 
         if with_spline:
             assert b_coeff.size != 0, "Beta spline coefficient array (t_coeff) is empty"
-            my_spline, ctr_pts, ctr_coeffs = _bspline(b_coeff)
-            self.my_spline = my_spline(self.s)
+            spline, ctr_pts, ctr_coeffs = _bspline(b_coeff)
+            self.spline = spline(self.s)
 
         else:
-
-            def constant_function(input):
-                """
-                Return array of ones same as the size of the input array. This
-                function is called when Beta spline function is not used.
-
-                Parameters
-                ----------
-                input
-
-                Returns
-                -------
-
-                """
-                return np.ones(input.shape)
-
-            self.my_spline = constant_function(self.s)
+            self.spline = np.full_like(self.s)
 
     def apply_torques(self, system, time: np.float64 = 0.0):
         self.compute_muscle_torques(
             time,
-            self.my_spline,
+            self.spline,
             self.s,
             self.angular_frequency,
             self.wave_number,
@@ -146,7 +125,7 @@ def apply_torques(self, system, time: np.float64 = 0.0):
     @njit(cache=True)
     def compute_muscle_torques(
         time,
-        my_spline,
+        spline,
         s,
         angular_frequency,
         wave_number,
@@ -167,7 +146,7 @@ def compute_muscle_torques(
             # front of wave number is positive, in Elastica cpp it is negative.
             torque_mag = (
                 factor
-                * my_spline
+                * spline
                 * np.sin(
                     angular_frequency * (time - switch_on_time - phase_shift)
                     - wave_number * s