diff --git a/elastica/joint.py b/elastica/joint.py index 082ae05a2..755b53bad 100644 --- a/elastica/joint.py +++ b/elastica/joint.py @@ -4,6 +4,7 @@ import numba from elastica.utils import Tolerance, MaxDimension from elastica._linalg import _batch_product_k_ik_to_ik +from elastica._rotations import _inv_rotate from math import sqrt @@ -209,9 +210,16 @@ class FixedJoint(FreeJoint): Damping coefficient of the joint. kt: float Rotational stiffness coefficient of the joint. + nut: float + Rotational damping coefficient of the joint. + rest_rotation_matrix: np.array + 2D (3,3) array containing data with 'float' type. + Rest 3x3 rotation matrix from system one to system two at the connected elements. + Instead of aligning the directors of both systems directly, a desired rest rotational matrix labeled C_12* + is enforced. """ - def __init__(self, k, nu, kt): + def __init__(self, k, nu, kt, nut=0.0, rest_rotation_matrix=None): """ Parameters @@ -222,51 +230,122 @@ def __init__(self, k, nu, kt): Damping coefficient of the joint. kt: float Rotational stiffness coefficient of the joint. + nut: float = 0. + Rotational damping coefficient of the joint. + rest_rotation_matrix: np.array + 2D (3,3) array containing data with 'float' type. + Rest 3x3 rotation matrix from system one to system two at the connected elements. + If provided, the rest rotation matrix is enforced between the two systems throughout the simulation. + If not provided, `rest_rotation_matrix` is initialized to the identity matrix, + which means that a restoring torque will be applied to align the directors of both systems directly. + (default=None) """ super().__init__(k, nu) # additional in-plane constraint through restoring torque # stiffness of the restoring constraint -- tuned empirically self.kt = kt + self.nut = nut + + # TODO: compute the rest rotation matrix directly during initialization + # as soon as systems (e.g. `rod_one` and `rod_two`) and indices (e.g. `index_one` and `index_two`) + # are available in the __init__ + if rest_rotation_matrix is None: + rest_rotation_matrix = np.eye(3) + assert rest_rotation_matrix.shape == (3, 3), "Rest rotation matrix must be 3x3" + self.rest_rotation_matrix = rest_rotation_matrix # Apply force is same as free joint def apply_forces(self, rod_one, index_one, rod_two, index_two): return super().apply_forces(rod_one, index_one, rod_two, index_two) - def apply_torques(self, rod_one, index_one, rod_two, index_two): - # current direction of the first element of link two - # also NOTE: - rod two is fixed at first element - link_direction = ( - rod_two.position_collection[..., index_two + 1] - - rod_two.position_collection[..., index_two] + def apply_torques(self, system_one, index_one, system_two, index_two): + # collect directors of systems one and two + # note that systems can be either rods or rigid bodies + system_one_director = system_one.director_collection[..., index_one] + system_two_director = system_two.director_collection[..., index_two] + + # rel_rot: C_12 = C_1I @ C_I2 + # C_12 is relative rotation matrix from system 1 to system 2 + # C_1I is the rotation from system 1 to the inertial frame (i.e. the world frame) + # C_I2 is the rotation from the inertial frame to system 2 frame (inverse of system_two_director) + rel_rot = system_one_director @ system_two_director.T + # error_rot: C_22* = C_21 @ C_12* + # C_22* is rotation matrix from current orientation of system 2 to desired orientation of system 2 + # C_21 is the inverse of C_12, which describes the relative (current) rotation from system 1 to system 2 + # C_12* is the desired rotation between systems one and two, which is saved in the static_rotation attribute + dev_rot = rel_rot.T @ self.rest_rotation_matrix + + # compute rotation vectors based on C_22* + # scipy implementation + # rot_vec = Rotation.from_matrix(dev_rot).as_rotvec() + # + # implementation using custom _inv_rotate compiled with numba + # rotation vector between identity matrix and C_22* + rot_vec = _inv_rotate(np.dstack([np.eye(3), dev_rot.T])).squeeze() + + # rotate rotation vector into inertial frame + rot_vec_inertial_frame = system_two_director.T @ rot_vec + + # deviation in rotation velocity between system 1 and system 2 + # first convert to inertial frame, then take differences + dev_omega = ( + system_two_director.T @ system_two.omega_collection[..., index_two] + - system_one_director.T @ system_one.omega_collection[..., index_one] ) - # To constrain the orientation of link two, the second node of link two should align with - # the direction of link one. Thus, we compute the desired position of the second node of link two - # as check1, and the current position of the second node of link two as check2. Check1 and check2 - # should overlap. + # we compute the constraining torque using a rotational spring - damper system in the inertial frame + torque = self.kt * rot_vec_inertial_frame - self.nut * dev_omega - tgt_destination = ( - rod_one.position_collection[..., index_one] - + rod_two.rest_lengths[index_two] * rod_one.tangents[..., index_one] - ) # dl of rod 2 can be different than rod 1 so use rest length of rod 2 + # The opposite torques will be applied to system one and two after rotating the torques into the local frame + system_one.external_torques[..., index_one] -= system_one_director @ torque + system_two.external_torques[..., index_two] += system_two_director @ torque - curr_destination = rod_two.position_collection[ - ..., index_two + 1 - ] # second element of rod2 - # Compute the restoring torque - forcedirection = -self.kt * ( - curr_destination - tgt_destination - ) # force direction is between rod2 2nd element and rod1 - torque = np.cross(link_direction, forcedirection) +def get_relative_rotation_two_systems(system_one, index_one, system_two, index_two): + """ + Compute the relative rotation matrix C_12 between system one and system two at the specified elements. - # The opposite torque will be applied on link one - rod_one.external_torques[..., index_one] -= ( - rod_one.director_collection[..., index_one] @ torque - ) - rod_two.external_torques[..., index_two] += ( - rod_two.director_collection[..., index_two] @ torque - ) + Examples + ---------- + How to get the relative rotation between two systems (e.g. the rotation from end of rod one to base of rod two): + + >>> rel_rot_mat = get_relative_rotation_two_systems(rod1, -1, rod2, 0) + + How to initialize a FixedJoint with a rest rotation between the two systems, + which is enforced throughout the simulation: + + >>> simulator.connect( + ... first_rod=rod1, second_rod=rod2, first_connect_idx=-1, second_connect_idx=0 + ... ).using( + ... FixedJoint, + ... ku=1e6, nu=0.0, kt=1e3, nut=0.0, + ... rest_rotation_matrix=get_relative_rotation_two_systems(rod1, -1, rod2, 0) + ... ) + + See Also + --------- + FixedJoint + + Parameters + ---------- + rod_one : object + Rod-like object + index_one : int + Index of first rod for joint. + rod_two : object + Rod-like object + index_two : int + Index of second rod for joint. + + Returns + ------- + relative_rotation_matrix : np.array + Relative rotation matrix C_12 between the two systems for their current state. + """ + return ( + system_one.director_collection[..., index_one] + @ system_two.director_collection[..., index_two].T + ) @numba.njit(cache=True) diff --git a/examples/JointCases/fixed_joint.py b/examples/JointCases/fixed_joint.py index 754d231cd..8c06c4270 100644 --- a/examples/JointCases/fixed_joint.py +++ b/examples/JointCases/fixed_joint.py @@ -104,36 +104,14 @@ class FixedJointSimulator( time_step=dt, ) -# Callback functions -# Add call backs -class TestJoints(CallBackBaseClass): - """ - Call back function for testing joints - """ - - 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["velocity"].append(system.velocity_collection.copy()) - return - - pp_list_rod1 = defaultdict(list) pp_list_rod2 = defaultdict(list) - fixed_joint_sim.collect_diagnostics(rod1).using( - TestJoints, step_skip=1000, callback_params=pp_list_rod1 + MyCallBack, step_skip=1000, callback_params=pp_list_rod1 ) fixed_joint_sim.collect_diagnostics(rod2).using( - TestJoints, step_skip=1000, callback_params=pp_list_rod2 + MyCallBack, step_skip=1000, callback_params=pp_list_rod2 ) fixed_joint_sim.finalize() diff --git a/examples/JointCases/fixed_joint_torsion.py b/examples/JointCases/fixed_joint_torsion.py new file mode 100644 index 000000000..fb13fa5cb --- /dev/null +++ b/examples/JointCases/fixed_joint_torsion.py @@ -0,0 +1,176 @@ +__doc__ = """Fixed joint example, for detailed explanation refer to Zhang et. al. Nature Comm. methods section.""" + +import matplotlib.pyplot as plt +import numpy as np +import sys + +# FIXME without appending sys.path make it more generic +sys.path.append("../../") +from elastica import * +from elastica.joint import get_relative_rotation_two_systems +from examples.JointCases.joint_cases_postprocessing import ( + plot_position, + plot_orientation, + plot_video, + plot_video_xy, + plot_video_xz, +) + + +class FixedJointSimulator( + BaseSystemCollection, Constraints, Connections, Forcing, Damping, CallBacks +): + pass + + +fixed_joint_sim = FixedJointSimulator() + +# setting up test params +n_elem = 10 +direction_rod1 = np.array([0.0, 0.0, 1.0]) +normal_rod1 = np.array([0.0, 1.0, 0.0]) +direction_rod2 = np.array([0.0, 1.0, 0.0]) +normal_rod2 = np.array([0.0, 0.0, 1.0]) +base_length = 0.2 +base_radius = 0.007 +base_area = np.pi * base_radius ** 2 +density = 1750 +E = 3e7 +poisson_ratio = 0.5 +shear_modulus = E / (poisson_ratio + 1.0) + +start_rod_1 = np.zeros((3,)) +start_rod_2 = start_rod_1 + direction_rod1 * base_length + +# Create rod 1 +rod1 = CosseratRod.straight_rod( + n_elem, + start_rod_1, + direction_rod1, + normal_rod1, + base_length, + base_radius, + density, + 0.0, # internal damping constant, deprecated in v0.3.0 + E, + shear_modulus=shear_modulus, +) +fixed_joint_sim.append(rod1) +# Create rod 2 +rod2 = CosseratRod.straight_rod( + n_elem, + start_rod_2, + direction_rod2, + normal_rod2, + base_length, + base_radius, + density, + 0.0, # internal damping constant, deprecated in v0.3.0 + E, + shear_modulus=shear_modulus, +) +fixed_joint_sim.append(rod2) + +# Apply boundary conditions to rod1. +fixed_joint_sim.constrain(rod1).using( + OneEndFixedBC, constrained_position_idx=(0,), constrained_director_idx=(0,) +) + +# Connect rod 1 and rod 2 +fixed_joint_sim.connect( + first_rod=rod1, second_rod=rod2, first_connect_idx=-1, second_connect_idx=0 +).using( + FixedJoint, + k=1e5, + nu=1.0, + kt=1e3, + nut=1e-3, + rest_rotation_matrix=get_relative_rotation_two_systems(rod1, -1, rod2, 0), +) + +# Add forces to rod2 +fixed_joint_sim.add_forcing_to(rod2).using( + UniformTorques, torque=5e-3, direction=np.array([0.0, 0.0, 1.0]) +) + +# add damping +damping_constant = 0.4 +dt = 1e-5 +fixed_joint_sim.dampen(rod1).using( + ExponentialDamper, + damping_constant=damping_constant, + time_step=dt, +) +fixed_joint_sim.dampen(rod2).using( + ExponentialDamper, + damping_constant=damping_constant, + time_step=dt, +) + + +pp_list_rod1 = defaultdict(list) +pp_list_rod2 = defaultdict(list) + + +fixed_joint_sim.collect_diagnostics(rod1).using( + MyCallBack, step_skip=1000, callback_params=pp_list_rod1 +) +fixed_joint_sim.collect_diagnostics(rod2).using( + MyCallBack, step_skip=1000, callback_params=pp_list_rod2 +) + +fixed_joint_sim.finalize() +timestepper = PositionVerlet() + +final_time = 10 +dl = base_length / n_elem +dt = 1e-5 +total_steps = int(final_time / dt) +print("Total steps", total_steps) +integrate(timestepper, fixed_joint_sim, final_time, total_steps) + + +plot_orientation( + "Orientation of last node of rod 1", + pp_list_rod1["time"], + np.array(pp_list_rod1["directors"])[..., -1], +) +plot_orientation( + "Orientation of last node of rod 2", + pp_list_rod2["time"], + np.array(pp_list_rod2["directors"])[..., -1], +) + +PLOT_FIGURE = True +SAVE_FIGURE = True +PLOT_VIDEO = True + +# plotting results +if PLOT_FIGURE: + filename = "fixed_joint_torsion_test.png" + plot_position(pp_list_rod1, pp_list_rod2, filename, SAVE_FIGURE) + +if PLOT_VIDEO: + filename = "fixed_joint_torsion_test" + fps = 100 # Hz + plot_video( + pp_list_rod1, + pp_list_rod2, + video_name=filename + ".mp4", + margin=0.2, + fps=fps, + ) + plot_video_xy( + pp_list_rod1, + pp_list_rod2, + video_name=filename + "_xy.mp4", + margin=0.2, + fps=fps, + ) + plot_video_xz( + pp_list_rod1, + pp_list_rod2, + video_name=filename + "_xz.mp4", + margin=0.2, + fps=fps, + ) diff --git a/examples/JointCases/hinge_joint.py b/examples/JointCases/hinge_joint.py index 416899e2b..f7b3bf904 100644 --- a/examples/JointCases/hinge_joint.py +++ b/examples/JointCases/hinge_joint.py @@ -106,39 +106,16 @@ class HingeJointSimulator( time_step=dt, ) -# Callback functions -# Add call backs -class TestJoints(CallBackBaseClass): - """ - Call back function for testing joints - """ - - 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["velocity"].append(system.velocity_collection.copy()) - return - - pp_list_rod1 = defaultdict(list) pp_list_rod2 = defaultdict(list) - hinge_joint_sim.collect_diagnostics(rod1).using( - TestJoints, step_skip=1000, callback_params=pp_list_rod1 + MyCallBack, step_skip=1000, callback_params=pp_list_rod1 ) hinge_joint_sim.collect_diagnostics(rod2).using( - TestJoints, step_skip=1000, callback_params=pp_list_rod2 + MyCallBack, step_skip=1000, callback_params=pp_list_rod2 ) - hinge_joint_sim.finalize() timestepper = PositionVerlet() # timestepper = PEFRL() diff --git a/examples/JointCases/joint_cases_postprocessing.py b/examples/JointCases/joint_cases_postprocessing.py index 87eb357b7..d7ea14376 100644 --- a/examples/JointCases/joint_cases_postprocessing.py +++ b/examples/JointCases/joint_cases_postprocessing.py @@ -2,6 +2,7 @@ from matplotlib import pyplot as plt from matplotlib.colors import to_rgb from mpl_toolkits import mplot3d +from scipy.spatial.transform import Rotation def plot_position( @@ -35,6 +36,25 @@ def plot_position( fig.savefig(filename) +def plot_orientation(title, time, directors): + quat = [] + for t in range(len(time)): + quat_t = Rotation.from_matrix(directors[t].T).as_quat() + quat.append(quat_t) + quat = np.array(quat) + + plt.figure(num=title) + plt.plot(time, quat[:, 0], label="x") + plt.plot(time, quat[:, 1], label="y") + plt.plot(time, quat[:, 2], label="z") + plt.plot(time, quat[:, 3], label="w") + plt.title(title) + plt.legend() + plt.xlabel("Time [s]") + plt.ylabel("Quaternion") + plt.show() + + def plot_video( plot_params_rod1: dict, plot_params_rod2: dict, @@ -94,7 +114,7 @@ def plot_video_xy( position_of_rod1 = np.array(plot_params_rod1["position"]) position_of_rod2 = np.array(plot_params_rod2["position"]) - print("plot video") + print("plot video xy") FFMpegWriter = manimation.writers["ffmpeg"] metadata = dict(title="Movie Test", artist="Matplotlib", comment="Movie support!") writer = FFMpegWriter(fps=fps, metadata=metadata) @@ -132,7 +152,7 @@ def plot_video_xz( position_of_rod1 = np.array(plot_params_rod1["position"]) position_of_rod2 = np.array(plot_params_rod2["position"]) - print("plot video") + print("plot video xz") FFMpegWriter = manimation.writers["ffmpeg"] metadata = dict(title="Movie Test", artist="Matplotlib", comment="Movie support!") writer = FFMpegWriter(fps=fps, metadata=metadata) diff --git a/examples/JointCases/spherical_joint.py b/examples/JointCases/spherical_joint.py index 54e5c7524..4164dd96c 100644 --- a/examples/JointCases/spherical_joint.py +++ b/examples/JointCases/spherical_joint.py @@ -107,37 +107,14 @@ class SphericalJointSimulator( time_step=dt, ) - -# Callback functions -# Add call backs -class TestJoints(CallBackBaseClass): - """ - Call back function for testing joints - """ - - 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["velocity"].append(system.velocity_collection.copy()) - return - - pp_list_rod1 = defaultdict(list) pp_list_rod2 = defaultdict(list) - spherical_joint_sim.collect_diagnostics(rod1).using( - TestJoints, step_skip=1000, callback_params=pp_list_rod1 + MyCallBack, step_skip=1000, callback_params=pp_list_rod1 ) spherical_joint_sim.collect_diagnostics(rod2).using( - TestJoints, step_skip=1000, callback_params=pp_list_rod2 + MyCallBack, step_skip=1000, callback_params=pp_list_rod2 ) spherical_joint_sim.finalize() diff --git a/tests/test_joint.py b/tests/test_joint.py index 042b32424..7c56e7372 100644 --- a/tests/test_joint.py +++ b/tests/test_joint.py @@ -6,8 +6,10 @@ from numpy.testing import assert_allclose from elastica.utils import Tolerance from elastica.rod.cosserat_rod import CosseratRod -import importlib import elastica +import importlib +import pytest +from scipy.spatial.transform import Rotation # TODO: change tests and made them independent of rod, at least assigin hardcoded values for forces and torques @@ -223,9 +225,19 @@ def test_hingejoint(): ) -def test_fixedjoint(): +rest_euler_angles = [ + np.array([0.0, 0.0, 0.0]), + np.array([np.pi / 2, 0.0, 0.0]), + np.array([0.0, np.pi / 2, 0.0]), + np.array([0.0, 0.0, np.pi / 2]), + 2 * np.pi * np.random.random_sample(size=3), +] + + +@pytest.mark.parametrize("rest_euler_angle", rest_euler_angles) +def test_fixedjoint(rest_euler_angle): # Define the rod for testing - # Some rod properties. We need them for constructer, they are not used. + # Some rod properties. We need them for constructor, they are not used. normal1 = np.array([0.0, 1.0, 0.0]) direction = np.array([1.0, 0.0, 0.0]) normal2 = np.array([0.0, 0.0, 1.0]) @@ -282,11 +294,18 @@ def test_fixedjoint(): rod1.velocity_collection = np.tile(v1, (1, n + 1)) rod2.velocity_collection = np.tile(v2, (1, n + 1)) - # Stiffness between points + # Rod angular velocity + omega1 = 1 / 180 * np.pi * np.array([0, 0, 1]).reshape(3, 1) + omega2 = -omega1 + rod1.omega_collection = np.tile(omega1, (1, n + 1)) + rod2.omega_collection = np.tile(omega2, (1, n + 1)) + + # Positional and rotational stiffness between systems k = 1e8 kt = 1e6 - # Damping between two points + # Positional and rotational damping between systems nu = 1 + nut = 1e2 # Rod indexes rod1_index = -1 @@ -309,7 +328,10 @@ def test_fixedjoint(): dampingforce = nu * normal_relative_vel * distance / end_distance contactforce = elasticforce - dampingforce - fxjt = FixedJoint(k, nu, kt) + rest_rotation_matrix = Rotation.from_euler( + "xyz", rest_euler_angle, degrees=False + ).as_matrix() + fxjt = FixedJoint(k, nu, kt, nut, rest_rotation_matrix=rest_rotation_matrix) fxjt.apply_forces(rod1, rod1_index, rod2, rod2_index) fxjt.apply_torques(rod1, rod1_index, rod2, rod2_index) @@ -321,28 +343,42 @@ def test_fixedjoint(): rod2.external_forces[..., rod2_index], -1 * contactforce, atol=Tolerance.atol() ) - linkdirection = ( - rod2.position_collection[..., rod2_index + 1] - - rod2.position_collection[..., rod2_index] - ) - - positiondiff = ( - rod1.position_collection[..., rod1_index] - - rod1.position_collection[..., rod1_index - 1] + # collect directors of systems one and two + # note that systems can be either rods or rigid bodies + rod1_director = rod1.director_collection[..., rod1_index] + rod2_director = rod2.director_collection[..., rod2_index] + + # rel_rot: C_12 = C_1I @ C_I2 + # C_12 is relative rotation matrix from system 1 to system 2 + # C_1I is the rotation from system 1 to the inertial frame (i.e. the world frame) + # C_I2 is the rotation from the inertial frame to system 2 frame (inverse of system_two_director) + rel_rot = rod1_director @ rod2_director.T + # error_rot: C_22* = C_21 @ C_12* + # C_22* is rotation matrix from current orientation of system 2 to desired orientation of system 2 + # C_21 is the inverse of C_12, which describes the relative (current) rotation from system 1 to system 2 + # C_12* is the desired rotation between systems one and two, which is saved in the static_rotation attribute + dev_rot = rel_rot.T @ rest_rotation_matrix + + # compute rotation vectors based on C_22* + # scipy implementation + rot_vec = Rotation.from_matrix(dev_rot).as_rotvec() + + # rotate rotation vector into inertial frame + rot_vec_inertial_frame = rod2_director.T @ rot_vec + + # deviation in rotation velocity between system 1 and system 2 + # first convert to inertial frame, then take differences + dev_omega = ( + rod2_director.T @ rod2.omega_collection[..., rod2_index] + - rod1_director.T @ rod1.omega_collection[..., rod1_index] ) - tangent = positiondiff / np.sqrt(np.dot(positiondiff, positiondiff)) - # rod 2 has to be alligned with rod 1 - check1 = ( - rod1.position_collection[..., rod1_index] - + rod2.rest_lengths[rod2_index] * tangent - ) - check2 = rod2.position_collection[..., rod2_index + 1] - forcedirection = -kt * (check2 - check1) - torque = np.cross(linkdirection, forcedirection) + # we compute the constraining torque using a rotational spring - damper system in the inertial frame + torque = kt * rot_vec_inertial_frame - nut * dev_omega - torque_rod1 = -rod1.director_collection[..., rod1_index] @ torque - torque_rod2 = rod2.director_collection[..., rod2_index] @ torque + # The opposite torques will be applied to system one and two after rotating the torques into the local frame + torque_rod1 = -rod1_director @ torque + torque_rod2 = rod2_director @ torque assert_allclose( rod1.external_torques[..., rod1_index], torque_rod1, atol=Tolerance.atol()