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] 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), + )