Update FixedJoints: restoring spring-damped-torques, initial rotation offset

elastica/joint.py

@@ -5,6 +5,7 @@
 from elastica.utils import Tolerance, MaxDimension
 from elastica._linalg import _batch_product_k_ik_to_ik
 from math import sqrt
+from scipy.spatial.transform import Rotation
 
 
 class FreeJoint:
@@ -209,9 +210,15 @@ 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
+            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 +229,96 @@ 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*
+        rot_vec = Rotation.from_matrix(dev_rot).as_rotvec()
+
+        # 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.
-
-        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
+        # 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
 
-        curr_destination = rod_two.position_collection[
-            ..., index_two + 1
-        ]  # second element of rod2
+        # 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
 
-        # 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)
 
-        # 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
-        )
+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.
+
+    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]
+    )
 
 
 @numba.njit(cache=True)

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()

examples/JointCases/fixed_joint_torsion.py

@@ -0,0 +1,167 @@
+__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 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 = np.array([0.0, 0.0, 1.0])
+normal = np.array([0.0, 1.0, 0.0])
+roll_direction = np.cross(direction, normal)
+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 * base_length
+
+# Create rod 1
+rod1 = CosseratRod.straight_rod(
+    n_elem,
+    start_rod_1,
+    direction,
+    normal,
+    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,
+    normal,
+    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)
+
+# 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,
+    )