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 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/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..fec1fc90 --- /dev/null +++ b/examples/ContinuumSnakeWithLiftingWaveCase/continuum_snake_with_lifting_wave.py @@ -0,0 +1,275 @@ +__doc__ = """Snake friction case from X. Zhang et. al. Nat. Comm. 2021""" + +import os +import numpy as np +import pickle + +from elastica import * + +from continuum_snake_postprocessing import ( + plot_snake_velocity, + plot_video, + compute_projected_velocity, + plot_curvature, +) +from snake_forcing import ( + MuscleTorquesLifting, +) +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, +): + # Initialize the simulation class + snake_sim = SnakeSimulator() + + # Simulation parameters + period = 2.0 + final_time = 20.0 + 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 = 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]) + 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 = 1e-1 + + # 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 = 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, + 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( + MuscleTorquesLifting, + 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, + ) + + # 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__": + # 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 + ) + + 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..68dea81d --- /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().__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..3b471aca --- /dev/null +++ b/examples/ContinuumSnakeWithLiftingWaveCase/snake_forcing.py @@ -0,0 +1,180 @@ +__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_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 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 + muscle torques as a traveling wave with a beta spline or only + as a traveling wave. For implementation details refer to X. Zhang et. al. Nat. Comm. 2021 + + 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. + switch_on_time: float + time to switch on the muscle activation. + """ + + def __init__( + self, + b_coeff, + period, + wave_number, + phase_shift, + direction, + rest_lengths, + ramp_up_time=0.0, + with_spline=False, + switch_on_time=0.0, + ): + """ + + Parameters + ---------- + 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. + switch_on_time: float + time to switch on the muscle activation. + + """ + super().__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 + + # 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" + spline, ctr_pts, ctr_coeffs = _bspline(b_coeff) + self.spline = spline(self.s) + + else: + self.spline = np.full_like(self.s) + + def apply_torques(self, system, time: np.float64 = 0.0): + self.compute_muscle_torques( + time, + self.spline, + self.s, + self.angular_frequency, + self.wave_number, + self.phase_shift, + self.ramp_up_time, + self.direction, + self.switch_on_time, + system.tangents, + system.director_collection, + system.external_torques, + ) + + @staticmethod + @njit(cache=True) + def compute_muscle_torques( + time, + spline, + s, + angular_frequency, + wave_number, + phase_shift, + ramp_up_time, + direction, + switch_on_time, + 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 + * 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. + # 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), + ) 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,