Reproducing snake diffraction results from your Nat Comm Paper

[AkashVardhan7]

Hello again, I am trying to develop a simulation of a worm passing through a pegged lattice, and started off with trying to replicate the results of snake diffraction and reflection and refraction reported in Fig 4 of your Nature communications paper.
I followed the instructions listed in the SI of your paper, and as a first pass tried to simulate a snake passing through 3 pegs located at (x=0, z = 0.75); (x=-0.23, z = 0.75) and (x=0.23, z = 0.75) where +z is the direction of snake travel.

The way I implemented this, is by modifying your continuumsnake.py example, and adding frictional forces corresponding to the patches, as shown in the following code snippet from the modified example.

Add friction forces

# Uniform friction with ground
origin_plane = np.array([0.0, -base_radius, 0.0])
normal_plane = normal
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]
static_mu_array = np.zeros(kinetic_mu_array.shape)
snake_diffraction_sim.add_forcing_to(shearable_rod).using(
    ea.AnisotropicFrictionalPlane,
    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,
)

# Add Patches at z = 0.75 m and several values of x
# Patch 1 -> 0,-base_radius,0.75
patch1_origin_plane = np.array([0.0,-base_radius,0.75])
patch1_normal_plane = normal
slip_velocity_tol = 1e-4
patch1_froude = 0.1
mu = base_length / (period * period * np.abs(gravitational_acc) * patch1_froude)
p_factor1 = 30
kinetic_mu_patch1_array = np.array (
    [mu*p_factor1, 1.5*mu*p_factor1, 2.0*mu*p_factor1]
)
static_mu_patch1_array = np.zeros(kinetic_mu_patch1_array.shape)
snake_diffraction_sim.add_forcing_to(shearable_rod).using(
    ea.AnisotropicFrictionalPlane,
    k=1.0,
    nu=1e-6,
    plane_origin=patch1_origin_plane,
    plane_normal=patch1_normal_plane,
    slip_velocity_tol=slip_velocity_tol,
    static_mu_array=static_mu_patch1_array,
    kinetic_mu_array=kinetic_mu_patch1_array,
)

# Patch 2 -> 0.023,-base_radius,0.75
patch2_origin_plane = np.array([0.023,-base_radius,0.75])
patch2_normal_plane = normal
slip_velocity_tol = 1e-4
patch2_froude = 0.1
mu = base_length / (period * period * np.abs(gravitational_acc) * patch2_froude)
p_factor2 = 30
kinetic_mu_patch2_array = np.array(
    [mu * p_factor2, 1.5 * mu * p_factor2, 2.0 * mu * p_factor2]
)
static_mu_patch2_array = np.zeros(kinetic_mu_patch1_array.shape)
snake_diffraction_sim.add_forcing_to(shearable_rod).using(
    ea.AnisotropicFrictionalPlane,
    k=1.0,
    nu=1e-6,
    plane_origin=patch2_origin_plane,
    plane_normal=patch2_normal_plane,
    slip_velocity_tol=slip_velocity_tol,
    static_mu_array=static_mu_patch2_array,
    kinetic_mu_array=kinetic_mu_patch2_array,
)

# Patch 3 -> -0.023,-base_radius,0.75
patch3_origin_plane = np.array([-0.023, -base_radius, 0.75])
patch3_normal_plane = normal
slip_velocity_tol = 1e-4
patch3_froude = 0.1
mu = base_length / (period * period * np.abs(gravitational_acc) * patch3_froude)
p_factor3 = 30
kinetic_mu_patch3_array = np.array(
    [mu * p_factor3, 1.5 * mu * p_factor3, 2.0 * mu * p_factor3]
)
static_mu_patch3_array = np.zeros(kinetic_mu_patch1_array.shape)
snake_diffraction_sim.add_forcing_to(shearable_rod).using(
    ea.AnisotropicFrictionalPlane,
    k=1.0,
    nu=1e-6,
    plane_origin=patch3_origin_plane,
    plane_normal=patch3_normal_plane,
    slip_velocity_tol=slip_velocity_tol,
    static_mu_array=static_mu_patch3_array,
    kinetic_mu_array=kinetic_mu_patch3_array,
)

I left the rest of the code unchanged, however while implementing this I had a question on defining the dimensions of the frictional patch which is supposed to represent a peg. The origin_plane which determines where the origin of the frictional substrate is going to be located only at that coordinate which we specify, however there is no way of specifying it's bounding dimensions such that when the snake is on any other part of the ground it should experience the standard frictional force, and only when it is in contact with the patch should it experience an increased friction.

Clearly the simulation doesn't work as intended and you can see it in the gif I am attaching. What should I do to implement the frictional patches correctly. Thanks.

[AkashVardhan7]

If I could find out how the ground element is created/referenced I would be able to partition into patches of higher friction based on the coordinates, but I can't find anything in your API which directly lets me extract coordinates of objects.

[xzhan139]

Hello Akash, I can help with this. Yes, your understanding is correct - the "AnisotropicFrictionalPlane" doesn't have a boundary, so by adding many of them together will just duplicate friction force.
The way to make local friction patches is to create a custom frictional plane. You can use the "AnisotropicFrictionalPlane" as a template, which should give you most of the basics you need. Then, the only change you want to make there, is to create a function to check the distance between your patch center and snake center. This sudo code may explain this:

scale_friction = 1
for every element/node on snake:
compute distance between this element to patch center
if distance < radius of the patch:
scale_friction = 30

This way, you can have different friction for each segment on the snake, and only add the friction when the element is in the patch region.
One more thing to add - in case you don't find the ground setup, it is in elastica/interaction.py
Please let me know if you have further question.

[AkashVardhan7]

Hi thanks for your reply @xzhan139 . I had a follow up question, when I try to create the customfrictional plane I would still need a way to bound the dimensions of the plane, where are the commands to set the boundaries inside the customfrictional plane?

[xzhan139]

Hi, let me explain this in another way. Your customfricitonal plane should start looking like a normal plane - no boundary and uniform friction everywhere. Then you just need to create a few high frictional spots on that plane. To do this, just scale the friction whenever snake elements are inside the region. So in the end, you just need one plane - your custom friction plane, which is an infinite plane, no bounds. Does this make sense?

[AkashVardhan7]

I get the idea, I haven't gotten the code to work yet.
Basically what I did was , create a function like you suggested which checks if a given point is inside a patch defined by known coordinates. And it works for test cases,

import numpy as np

def is_point_inside_patch(point, patch_vertices):
"""
Check if a 3D point is inside a patch defined by four vertices.

Args:
    point (numpy.ndarray): 3D coordinates of the point.
    patch_vertices (list of numpy.ndarray): List of four 3D coordinates defining the patch.

Returns:
    bool: True if the point is inside the patch, False otherwise.
"""
# Calculate the normal vectors of the four triangles formed by the point and patch vertices
normals = []
for i in range(4):
    vertex1 = patch_vertices[i]
    vertex2 = patch_vertices[(i + 1) % 4]
    normal = np.cross(vertex2 - vertex1, point - vertex1)
    normals.append(normal)

# Check if the point is on the same side of all the patch's triangles (using dot product)
for i in range(4):
    if np.dot(normals[i], normals[(i + 1) % 4]) < 0:
        return False

return True

Test cases

if name == "main":
# Define patch vertices (four known points in 3D space)
patch_vertices = [
np.array([0.0, 0.0, 0.0]),
np.array([0.5, 0.0, 0.0]),
np.array([0.5, 0.5, 0.0]),
np.array([0.0, 0.5, 0.0]),
]

# Generate 10 random points
random_points = np.random.uniform(size=(10, 3))

# Check if each random point is inside the patch and print the result
for i, point in enumerate(random_points):
    is_inside = is_point_inside_patch(point, patch_vertices)
    print(f"Point {i + 1}: {point} is inside the patch: {is_inside}")
    
    
    
    
    
    Now when I do the same thing in the main sim, snake_diffraction.py. Which I am copy pasting... I am running into an error.
    
    
    Here's the code.
    import os

import numpy as np
import elastica as ea

def is_point_inside_patch(point, patch_vertices):
"""
Check if a 3D point is inside a patch defined by four vertices.

Args:
    point (numpy.ndarray): 3D coordinates of the point.
    patch_vertices (list of numpy.ndarray): List of four 3D coordinates defining the patch.

Returns:
    bool: True if the point is inside the patch, False otherwise.
"""
# Calculate the normal vectors of the four triangles formed by the point and patch vertices
normals = []
for i in range(4):
    vertex1 = patch_vertices[i]
    vertex2 = patch_vertices[(i + 1) % 4]
    normal = np.cross(vertex2 - vertex1, point - vertex1)
    normals.append(normal)

# Check if the point is on the same side of all the patch's triangles (using dot product)
for i in range(4):
    if np.dot(normals[i], normals[(i + 1) % 4]) < 0:
        return False

return True

from snake_diffraction_postprocessing import (
plot_snake_velocity,
plot_video,
compute_projected_velocity,
plot_curvature,
)

class SnakeDiffractionSimulator(
ea.BaseSystemCollection, ea.Constraints, ea.Forcing, ea.Damping, ea.CallBacks
):
pass

def run_snake(
b_coeff, PLOT_FIGURE=False, SAVE_FIGURE=False, SAVE_VIDEO=False, SAVE_RESULTS=False
):
# Initialize the simulation class
snake_diffraction_sim = SnakeDiffractionSimulator()

# Simulation parameters
period = 2
final_time = (20.0 + 0.01) * period

# setting up test params
n_elem = 50
start = np.zeros((3,))
direction = np.array([0.0, 0.0, 1.0])
normal = np.array([0.0, 1.0, 0.0])
base_length = 0.35
base_radius = base_length * 0.011
density = 1000
E = 1e6
poisson_ratio = 0.5
shear_modulus = E / (poisson_ratio + 1.0)

snake_body = ea.CosseratRod.straight_rod(
    n_elem,
    start,
    direction,
    normal,
    base_length,
    base_radius,
    density,
    youngs_modulus=E,
    shear_modulus=shear_modulus,
)

snake_diffraction_sim.append(snake_body)

snake_positions = snake_body.position_collection
print("snake position is...")
print(snake_positions)

# Add gravitational forces
gravitational_acc = -9.80665
snake_diffraction_sim.add_forcing_to(snake_body).using(
    ea.GravityForces, acc_gravity=np.array([0.0, gravitational_acc, 0.0])
)

# Add muscle torques only Lateral Wave right now
wave_length = b_coeff[-1]
snake_diffraction_sim.add_forcing_to(snake_body).using(
    ea.MuscleTorques,
    base_length=base_length,
    b_coeff=b_coeff[:-1],
    period=period,
    wave_number=2.0 * np.pi / (wave_length),
    phase_shift=0.0,
    rest_lengths=snake_body.rest_lengths,
    ramp_up_time=period,
    direction=normal,
    with_spline=True,
)

# Add friction forces
# Uniform friction with ground
origin_plane = np.array([0.0, -base_radius, 0.0])
normal_plane = normal
slip_velocity_tol_Regular = 1e-6
slip_velocity_tol_HighFriction = 1e-4
p_factor = 30
froude = 0.1
mu = base_length / (period * period * np.abs(gravitational_acc) * froude)
kinetic_mu_array_Regular = np.array(
    [mu, 2 * mu, 10.0 * mu]
)  # [forward, backward, sideways]

static_mu_array_Regular = np.zeros(kinetic_mu_array_Regular.shape)

kinetic_mu_array_HighFriction = np.array(
    [mu * p_factor, 2 * mu * p_factor, 10.0 * mu * p_factor]
)

static_mu_array_HighFriction = np.zeros(kinetic_mu_array_HighFriction.shape)

patch1_A_coord = np.array([0.15, 0.0, 0.6])
patch1_B_coord = np.array([0.15, 0.0, 0.9])
patch1_C_coord = np.array([-0.15, 0.0, 0.9])
patch1_D_coord = np.array([-0.15, 0.0, 0.6])

patch_vertices = [
    patch1_A_coord,
    patch1_B_coord,
    patch1_C_coord,
    patch1_D_coord,
]

for i, node_position in enumerate(snake_body.position_collection):
    # Check if the node position is inside the patch
    if is_point_inside_patch(node_position, patch_vertices):
        # Apply high friction coefficients
        snake_diffraction_sim.add_forcing_to(snake_body).using(
            ea.AnisotropicFrictionalPlane,
            k=1.0,  # Modify this to your desired value
            nu=1e-6,  # Modify this to your desired value
            plane_origin=origin_plane,
            plane_normal=normal_plane,
            slip_velocity_tol=slip_velocity_tol_HighFriction,  # Define this value
            static_mu_array=static_mu_array_HighFriction,  # Define this array
            kinetic_mu_array=kinetic_mu_array_HighFriction,  # Define this array
        )
    else:
        # Apply regular friction coefficients
        snake_diffraction_sim.add_forcing_to(snake_body).using(
            ea.AnisotropicFrictionalPlane,
            k=1.0,  # Modify this to your desired value
            nu=1e-6,  # Modify this to your desired value
            plane_origin=origin_plane,
            plane_normal=normal_plane,
            slip_velocity_tol=slip_velocity_tol_Regular,
            static_mu_array=static_mu_array_Regular,
            kinetic_mu_array=kinetic_mu_array_Regular,
        )

# snake_diffraction_sim.add_forcing_to(snake_body).using(
#    ea.AnisotropicFrictionalPlane,
#    k=1.0,
#    nu=1e-6,
#    plane_origin=origin_plane,
#    plane_normal=normal_plane,
#   slip_velocity_tol=slip_velocity_tol_LessFriction,
#    static_mu_array=static_mu_array_LessFriction,
#    kinetic_mu_array=kinetic_mu_array_LessFriction,
# )

# add damping
damping_constant = 2e-3
time_step = 1e-4
snake_diffraction_sim.dampen(snake_body).using(
    ea.AnalyticalLinearDamper,
    damping_constant=damping_constant,
    time_step=time_step,
)

# Implement Pegs Via Collisions
# start = np.array([[0.0, -base_radius, 0.75]])
# Peg_n_elem = 50
# start_peg1 = start
# direction_peg1 = np.array([0.0, 1.0, 0.0])
# normal_peg1 = np.array([0.0, 0.0, 1.0])

# Rod parameters
# Peg_base_length = 0.05
# Peg_base_radius = 0.023
# Peg_base_area = np.pi * base_radius ** 2
# Peg_density = 1750
# Peg_nu = 0.0
# Peg_E = 5e6
# Peg_poisson_ratio = 0.5
# eg_shear_modulus = Peg_E / (Peg_poisson_ratio + 1.0)

# peg1 = ea.CosseratRod.straight_rod(
#    Peg_n_elem,
#    start_peg1,
#    direction_peg1,
#   normal_peg1,
#    Peg_base_length,
#    Peg_base_radius,
#    Peg_density,
#    youngs_modulus=Peg_E,
#    shear_modulus=Peg_shear_modulus,
# )

# snake_diffraction_sim.append(peg1)

# Contact between two rods
# snake_diffraction_sim.connect(shearable_rod, peg1).using(
#   ea.ExternalContact, k=1e3, nu=0.5
# )

total_steps = int(final_time / time_step)
rendering_fps = 60
step_skip = int(1.0 / (rendering_fps * time_step))

# Add call backs
class ContinuumSnakeCallBack(ea.CallBackBaseClass):
    """
    Call back function for continuum snake
    """

    def __init__(self, step_skip: int, callback_params: dict):
        ea.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["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 = ea.defaultdict(list)
snake_diffraction_sim.collect_diagnostics(snake_body).using(
    ContinuumSnakeCallBack, step_skip=step_skip, callback_params=pp_list
)

snake_diffraction_sim.finalize()

timestepper = ea.PositionVerlet()
ea.integrate(timestepper, snake_diffraction_sim, final_time, total_steps)

if PLOT_FIGURE:
    filename_plot = "snake_diffraction_velocity2.png"
    plot_snake_velocity(pp_list, period, filename_plot, SAVE_FIGURE)
    plot_curvature(pp_list, snake_body.rest_lengths, period, SAVE_FIGURE)

    if SAVE_VIDEO:
        filename_video = 'snake_diffraction2.gif'
        # filename_video = "continuum_snake.mp4"
        plot_video(
            pp_list,
            video_name=filename_video,
            fps=rendering_fps,
            xlim=(0, 4),
            ylim=(-1, 1),
        )

if SAVE_RESULTS:
    import pickle

    filename = "snake_diffraction.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":

# Options
PLOT_FIGURE = True
SAVE_FIGURE = True
SAVE_VIDEO = True
SAVE_RESULTS = False
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(
            [3.4e-3, 3.3e-3, 4.2e-3, 2.6e-3, 3.6e-3, 3.5e-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
    )

    print("average forward velocity:", avg_forward)
    print("average forward lateral:", avg_lateral)

And here's the error...

Traceback (most recent call last):
File "C:\Users\avardhan7\AppData\Local\JetBrains\PyCharm 2023.1.2\plugins\python\helpers\pydev\pydevconsole.py", line 364, in runcode
coro = func()
File "", line 1, in
File "C:\Users\avardhan7\AppData\Local\JetBrains\PyCharm 2023.1.2\plugins\python\helpers\pydev_pydev_bundle\pydev_umd.py", line 198, in runfile
pydev_imports.execfile(filename, global_vars, local_vars) # execute the script
File "C:\Users\avardhan7\AppData\Local\JetBrains\PyCharm 2023.1.2\plugins\python\helpers\pydev_pydev_imps_pydev_execfile.py", line 18, in execfile
exec(compile(contents+"\n", file, 'exec'), glob, loc)
File "C:\Users\avardhan7\Documents\OPEN_AI\pyElastica\snake_diffraction\snake_diffraction.py", line 351, in
[avg_forward, avg_lateral, pp_list] = run_snake(
File "C:\Users\avardhan7\Documents\OPEN_AI\pyElastica\snake_diffraction\snake_diffraction.py", line 143, in run_snake
if is_point_inside_patch(node_position, patch_vertices):
File "C:\Users\avardhan7\Documents\OPEN_AI\pyElastica\snake_diffraction\snake_diffraction.py", line 22, in is_point_inside_patch
normal = np.cross(vertex2 - vertex1, point - vertex1)
ValueError: operands could not be broadcast together with shapes (51,) (3,)

I haven't had time to debug this...you can choose to keep this thread open or close it, and I will post with the status once I am able to resolve the issue. Thanks.

[AkashVardhan7]

Finally got around to getting it to work, there was a minor error with the loop syntax.
Here's a video. Thanks for your help.

[armantekinalp]

Seems like this issue has been resolved, I am closing it.