Using SeqIKPy for Drosophila locomotion#
In this tutorial, we will use seqikpy to perform inverse kinematics on a locomotion data.
# Import the necessary libraries
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from seqikpy.utils import load_file, save_file, calculate_body_size, dict_to_nparray_pose
from seqikpy.alignment import AlignPose, convert_from_df3dpp_to_dict
from seqikpy.kinematic_chain import KinematicChainSeq
from seqikpy.leg_inverse_kinematics import LegInvKinSeq
from seqikpy.visualization import plot_3d_points, animate_3d_points
# Set up the constant variables
leg_joint_angle_names = [
"ThC_yaw",
"ThC_pitch",
"ThC_roll",
"CTr_pitch",
"CTr_roll",
"FTi_pitch",
"TiTa_pitch",
]
legs_to_align = ["RF", "RM", "RH", "LF", "LM", "LH"]
# Load the data, below contains the kinematics during locomotion
data_path = Path("../data/df3d_pose_result__210902_PR_Fly1")
pose_data = load_file(
data_path / "pose_result__210902_PR_Fly1_aligned.pkl"
# Kinematics during grooming
# "../data/df3d_180921_aDN_CsCh_Fly6_003_SG1/aligned_pose__180921_aDN_CsCh_Fly6_003_SG1_behData_images_images.pkl"
)
# Convert the dictionary in a compatible format
converted_dict_full = convert_from_df3dpp_to_dict(pose_data)
# Select the time period where locomotion happens
converted_dict = {}
for key, dict_v in converted_dict_full.items():
converted_dict[key] = dict_v[300:400,:,:]
Defining the body model#
It is very challenging to precisely define the joint rotational axes in Drosophila due to the complex structure of their exoskeleton. Here, we make use of a previously published biomechanical model of Drosophila melonagestor, NeuroMechFly. Specifically, we identified the 3D points corresponding to each joint on the biomechanical body, and use this information to transform the real animal 3D data into the biomechanical modelās frame of reference.
Biomechanical model
TEMPLATE_NMF_LOCOMOTION = {
"RF_Coxa": np.array([0.35, -0.27, 0.400]),
"RF_Femur": np.array([0.35, -0.27, -0.025]),
"RF_Tibia": np.array([0.35, -0.27, -0.731]),
"RF_Tarsus": np.array([0.35, -0.27, -1.249]),
"RF_Claw": np.array([0.35, -0.27, -1.912]),
"LF_Coxa": np.array([0.35, 0.27, 0.400]),
"LF_Femur": np.array([0.35, 0.27, -0.025]),
"LF_Tibia": np.array([0.35, 0.27, -0.731]),
"LF_Tarsus": np.array([0.35, 0.27, -1.249]),
"LF_Claw": np.array([0.35, 0.27, -1.912]),
"RM_Coxa": np.array([0, -0.125, 0]),
"RM_Femur": np.array([0, -0.125, -0.182]),
"RM_Tibia": np.array([0, -0.125, -0.965]),
"RM_Tarsus": np.array([0, -0.125, -1.633]),
"RM_Claw": np.array([0, -0.125, -2.328]),
"LM_Coxa": np.array([0, 0.125, 0]),
"LM_Femur": np.array([0, 0.125, -0.182]),
"LM_Tibia": np.array([0, 0.125, -0.965]),
"LM_Tarsus": np.array([0, 0.125, -1.633]),
"LM_Claw": np.array([0, 0.125, -2.328]),
"RH_Coxa": np.array([-0.215, -0.087, -0.073]),
"RH_Femur": np.array([-0.215, -0.087, -0.272]),
"RH_Tibia": np.array([-0.215, -0.087, -1.108]),
"RH_Tarsus": np.array([-0.215, -0.087, -1.793]),
"RH_Claw": np.array([-0.215, -0.087, -2.588]),
"LH_Coxa": np.array([-0.215, 0.087, -0.073]),
"LH_Femur": np.array([-0.215, 0.087, -0.272]),
"LH_Tibia": np.array([-0.215, 0.087, -1.108]),
"LH_Tarsus": np.array([-0.215, 0.087, -1.793]),
"LH_Claw": np.array([-0.215, 0.087, -2.588]),
}
Aligning the 3D pose data#
The alignment process serves two purposes:
Transforming the 3D pose data to the biomechanical modelās frame of reference.
(Optional) scaling the animal 3D data to match the size of the biomechanical model if the biomechanical modelās size is given to the
AlignPoseclass. This option is useful if you want to scale all of your data to one body size, or if you want to perform kinematic replay in simulation using the body model. Alternatively, the user can provide the experimental animalās body size into the class. In that case, the scaling will ensure that the animalās limb size is the same throughout the trial duration.
# lets align the pose
align = AlignPose(
pose_data_dict=converted_dict,
legs_list=legs_to_align,
include_claw=False,
body_template=TEMPLATE_NMF_LOCOMOTION,
# if body_size is none, then the size will be
# calculated from the template
body_size=None,
log_level="INFO"
)
aligned_pos = align.align_pose(export_path=data_path)
2024-02-14 16:04:43,609 - INFO- Scale factor for RF leg: 1.006482293211129
2024-02-14 16:04:43,612 - INFO- Scale factor for RM leg: 1.026667148098488
2024-02-14 16:04:43,616 - INFO- Scale factor for RH leg: 1.0121829552693662
2024-02-14 16:04:43,620 - INFO- Scale factor for LF leg: 1.0138788147880813
2024-02-14 16:04:43,624 - INFO- Scale factor for LM leg: 1.0134311442162325
2024-02-14 16:04:43,628 - INFO- Scale factor for LH leg: 1.0027376149059155
2024-02-14 16:04:43,630 - INFO- Aligned pose is saved at ../data/df3d_pose_result__210902_PR_Fly1
Sequential Inverse Kinematics#
Below, we will start the inverse kinematics process using seqikpy
The initial step is to determine the initial seeds for the first time step. While it is not crucial to have good initial seeds, having initial values out of the joint DOF limits will raise an error.
# We are defining a set of initial seeds for the IK optimization
INITIAL_ANGLES_LOCOMOTION = {
"RF": {
# Base ThC yaw pitch CTr pitch
"stage_1": np.array([0.0, 0.45, -0.07, -2.14]),
# BaseĀ ThC yaw pitch roll CTr pitch CTr roll
"stage_2": np.array([0.0, 0.45, -0.07, -0.32, -2.14, 1.4]),
# Base ThC yaw pitch roll CTr pitch CTr roll FTi pitch
"stage_3": np.array([0.0, 0.45, -0.07, -0.32, -2.14, -1.25, 1.48, 0.0]),
# Base ThC yaw pitch roll CTr pitch CTr roll FTi pitch TiTa pitch
"stage_4": np.array([0.0, 0.45, -0.07, -0.32, -2.14, -1.25, 1.48, 0.0, 0.0]),
},
"LF": {
"stage_1": np.array([0.0, -0.45, -0.07, -2.14]),
"stage_2": np.array([0.0, -0.45, -0.07, 0.32, -2.14, 1.4]),
"stage_3": np.array([0.0, -0.45, -0.07, 0.32, -2.14, 1.25, 1.48, 0.0]),
"stage_4": np.array([0.0, -0.45, -0.07, 0.32, -2.14, 1.25, 1.48, 0.0, 0.0]),
},
"RM": {
"stage_1": np.array([0.0, 0.45, 0.37, -2.14]),
"stage_2": np.array([0.0, 0.45, 0.37, -0.32, -2.14, 1.4]),
"stage_3": np.array([0.0, 0.45, 0.37, -0.32, -2.14, -1.25, 1.48, 0.0]),
"stage_4": np.array([0.0, 0.45, 0.37, -0.32, -2.14, -1.25, 1.48, 0.0, 0.0]),
},
"LM": {
"stage_1": np.array([0.0, -0.45, 0.37, -2.14]),
"stage_2": np.array([0.0, -0.45, 0.37, 0.32, -2.14, 1.4]),
"stage_3": np.array([0.0, -0.45, 0.37, 0.32, -2.14, 1.25, 1.48, 0.0]),
"stage_4": np.array([0.0, -0.45, 0.37, 0.32, -2.14, 1.25, 1.48, 0.0, 0.0]),
},
"RH": {
"stage_1": np.array([0.0, 0.45, 0.07, -2.14]),
"stage_2": np.array([0.0, 0.45, 0.07, -0.32, -2.14, 1.4]),
"stage_3": np.array([0.0, 0.45, 0.07, -0.32, -2.14, -1.25, 1.48, 0.0]),
"stage_4": np.array([0.0, 0.45, 0.07, -0.32, -2.14, -1.25, 1.48, 0.0, 0.0]),
},
"LH": {
"stage_1": np.array([0.0, -0.45, 0.07, -2.14]),
"stage_2": np.array([0.0, -0.45, 0.07, 0.32, -2.14, 1.4]),
"stage_3": np.array([0.0, -0.45, 0.07, 0.32, -2.14, 1.25, 1.48, 0.0]),
"stage_4": np.array([0.0, -0.45, 0.07, 0.32, -2.14, 1.25, 1.48, 0.0, 0.0]),
},
}
Secondly, we need to define some lower and upper limits for the joints. Again, having informative bounds is helpful to prevent singularities in the optimization. For example, Tibia-tarsus pitch joint should not exceed 0 degrees (always negative since the rotation is clockwise). So, it makes sense to assign an upper limit of 0 degrees for this joint.
# We determine the bounds for each joint DOF
BOUNDS_LOCOMOTION = {
"RF_ThC_yaw": (-3.141592653589793, 3.141592653589793),
"RF_ThC_pitch": (np.deg2rad(-90), np.deg2rad(90)),
"RF_ThC_roll": (-3.141592653589793, 3.141592653589793),
"RF_CTr_pitch": (-3.141592653589793, 3.141592653589793),
"RF_FTi_pitch": (-3.141592653589793, 3.141592653589793),
"RF_CTr_roll": (-3.141592653589793, 3.141592653589793),
"RF_TiTa_pitch": (-3.141592653589793, np.deg2rad(0)),
"RM_ThC_yaw": (np.deg2rad(-50), np.deg2rad(50)),
"RM_ThC_pitch": (-3.141592653589793, 3.141592653589793),
"RM_ThC_roll": (-3.141592653589793, 0),
"RM_CTr_pitch": (-3.141592653589793, 3.141592653589793),
"RM_FTi_pitch": (-3.141592653589793, 3.141592653589793),
"RM_CTr_roll": (-3.141592653589793, 3.141592653589793),
"RM_TiTa_pitch": (-3.141592653589793, np.deg2rad(0)),
"RH_ThC_yaw": (np.deg2rad(-50), np.deg2rad(50)),
"RH_ThC_pitch": (np.deg2rad(-50), np.deg2rad(50)),
"RH_ThC_roll": (-3.141592653589793, 0),
"RH_CTr_pitch": (np.deg2rad(-180), np.deg2rad(0)),
"RH_FTi_pitch": (-3.141592653589793, 3.141592653589793),
"RH_CTr_roll": (-3.141592653589793, 3.141592653589793),
"RH_TiTa_pitch": (-3.141592653589793, np.deg2rad(0)),
"LF_ThC_yaw": (-3.141592653589793, 3.141592653589793),
"LF_ThC_pitch": (np.deg2rad(-90), np.deg2rad(90)),
"LF_ThC_roll": (-3.141592653589793, 3.141592653589793),
"LF_CTr_pitch": (-3.141592653589793, 3.141592653589793),
"LF_FTi_pitch": (-3.141592653589793, 3.141592653589793),
"LF_CTr_roll": (-3.141592653589793, 3.141592653589793),
"LF_TiTa_pitch": (-3.141592653589793, np.deg2rad(0)),
"LM_ThC_yaw": (np.deg2rad(-50), np.deg2rad(50)),
"LM_ThC_pitch": (-3.141592653589793, 3.141592653589793),
"LM_ThC_roll": (0, 3.141592653589793),
"LM_CTr_pitch": (-3.141592653589793, 3.141592653589793),
"LM_FTi_pitch": (-3.141592653589793, 3.141592653589793),
"LM_CTr_roll": (-3.141592653589793, 3.141592653589793),
"LM_TiTa_pitch": (-3.141592653589793, np.deg2rad(0)),
"LH_ThC_yaw": (np.deg2rad(-50), np.deg2rad(50)),
"LH_ThC_pitch": (np.deg2rad(-50), np.deg2rad(50)),
"LH_ThC_roll": (0, 3.141592653589793),
"LH_CTr_pitch": (np.deg2rad(-180), np.deg2rad(0)),
"LH_FTi_pitch": (-3.141592653589793, 3.141592653589793),
"LH_CTr_roll": (-3.141592653589793, 3.141592653589793),
"LH_TiTa_pitch": (-3.141592653589793, np.deg2rad(0)),
}
Run sequential inverse kinematics#
# Initialize the necessary classes
kin_chain = KinematicChainSeq(
bounds_dof=BOUNDS_LOCOMOTION,
body_size=calculate_body_size(
TEMPLATE_NMF_LOCOMOTION,
legs_to_align
),
# perform ik on all 6 legs
legs_list=legs_to_align,
)
class_seq_ik = LegInvKinSeq(
aligned_pos=aligned_pos,
kinematic_chain_class=kin_chain,
initial_angles=INITIAL_ANGLES_LOCOMOTION,
log_level="INFO",
)
The below code will commence the IK process, which consists of four stages per kinematic chain. For more details, please visit the methodology page.
After the process is over, run_ik_and_fk will save the results in two separete .pkl files:
leg_joint_angles.pkl-> pickle file that contains the leg joint anglesforward_kinematics.pkl-> 3D position of the legs construction from the calculated leg joint angles, useful for visualization and debugging purposes
leg_joint_angles, forward_kinematics = class_seq_ik.run_ik_and_fk(
export_path=data_path,
hide_progress_bar=True
)
2024-02-14 16:04:43,763 - INFO- Computing joint angles and forward kinematics...
2024-02-14 16:06:24,251 - INFO- Joint angles and forward kinematics are saved at ../data/df3d_pose_result__210902_PR_Fly1
2D visualization of the joint angles#
# Let's plot the joint angles for all six legs
fig, axs = plt.subplots(3, 2, figsize=(9, 6), dpi=200)
axs = axs.flatten()
for angle_name in leg_joint_angle_names:
for i, leg_name in enumerate(["RF", "LF", "RM", "LM", "RH", "LH"]):
axs[i].plot(
np.rad2deg(leg_joint_angles[f"Angle_{leg_name}_{angle_name}"]),
label=angle_name,
lw=2,
)
axs[i].set_ylabel(leg_name)
time_step = 1e-2
for ax in axs:
ax.set_xticklabels(np.array(ax.get_xticks() * time_step, dtype='f'))
axs[-1].set_xlabel("Time (sec)")
axs[-2].set_xlabel("Time (sec)")
axs[1].legend(bbox_to_anchor=(1.1, 1), frameon=False)
plt.suptitle("Leg joint angles (deg)")
plt.tight_layout()
# fig.savefig("../results/alljoints_follow_ikpy_left.png", bbox_inches="tight")
plt.show()
Animation of the target 3D pose and the forward kinematics#
#Ā Files saved automatically after the IK-FK process above
azim=0
elev=90
animate_3d_points(
points3d=aligned_pos,
points3d_second=forward_kinematics,
export_path=data_path / f'fk_ik_elev_{elev}_azim_{azim}.mp4',
frame_no=100,
elev=elev,
azim=azim,
fps=100,
)
2024-02-14 16:08:32,081 - INFO- Making animation...
2024-02-14 16:08:35,692 - INFO- Animation is saved at ../data/df3d_pose_result__210902_PR_Fly1/fk_ik_elev_90_azim_0.mp4
Original 3D pose (solid) F.K. (dashed)#
# If wanted, you can save each frame as a png file
plt.style.use('default')
t = 50
for azim, elev in [(0,0),(90,0),(0,90)]:
fig = plt.figure(dpi=100)
ax3d = fig.add_subplot(projection='3d')
ax3d.view_init(azim=azim, elev=elev)
ax3d.set_xlabel('x')
ax3d.set_ylabel('y')
ax3d.set_zlabel('z')
plot_3d_points(
ax3d,
aligned_pos,
t=t,
line_style='solid'
)
ax3d.legend(bbox_to_anchor=(1.3,0.6))
plot_3d_points(
ax3d,
forward_kinematics,
t=t,
line_style='--'
)
ax3d.set_title('Solid - raw 3D, Dashed - FK', y=0.95)
plt.show()
