API Reference

Submodules

seqikpy.alignment module

Code for aligning 3D pose to a fly template. The best practice for getting good alignment is to have an accurate 3D pose and a template whose key points are matching the tracked key points closely.

This class expects a 3D file in the following format

>>> pose_data_dict = {
        "<side (R,L)><segment (F, M, H)>_leg": np.ndarray[N_frames,5,3],
        "<side (R,L)>_head": np.ndarray[N_frames,2,3],
        "Neck": np.ndarray[N_frames,1,3],
    }

Usage might differ based on the needs. For now, there are three cases that you can use the class with:

Case 1: we have 3D pose obtained, and we would like to align it but first we need to convert the pose data into a dictionary format NOTE: if the 3D pose is not in the format described above, then you need to * Convert your 3D pose file manually to the required format * Or, if you obtain the 3D pose from anipose, simply set convert_func to convert_from_anipose_to_dict .

>>> data_path = Path("../data/anipose_220525_aJO_Fly001_001/pose-3d")
>>> align = AlignPose.from_file_path(
>>>     main_dir=data_path,
>>>     file_name="pose3d.h5",
>>>     legs_list=["RF","LF"],
>>>     convert_func=convert_from_anipose_to_dict,
>>>     pts2align=PTS2ALIGN,
>>>     include_claw=False,
>>>     body_template=NMF_TEMPLATE,
>>> )
>>> aligned_pos = align.align_pose(export_path=data_path)

Case 2: we have a pose data in the required data structure, we just want to load and align it

>>> data_path = Path("../data/anipose_220525_aJO_Fly001_001/pose-3d")
>>> align = AlignPose.from_file_path(
>>>     main_dir=data_path,
>>>     file_name="converted_pose_dict.pkl",
>>>     legs_list=["RF","LF"],
>>>     convert_func=None,
>>>     pts2align=PTS2ALIGN,
>>>     include_claw=False,
>>>     body_template=NMF_TEMPLATE,
>>> )
>>> aligned_pos = align.align_pose(export_path=data_path)

Case 3: we have a pose data in the required format loaded already, we want to feed it to the class and align the pose. This assumes that the pose data is already aligned in the right format. If not, use the static method convert_from_anipose.

>>> data_path = Path("../data/anipose_220525_aJO_Fly001_001/pose-3d")
>>> f_path = data_path / "converted_pose_dict.pkl"
>>> with open(f_path, "rb") as f:
>>>     pose_data = pickle.load(f)
>>> align = AlignPose(
>>>     pose_data_dict=pose_data,
>>>     legs_list=["RF","LF"],
>>>     include_claw=False,
>>>     body_template=NMF_TEMPLATE,
>>> )
>>> aligned_pos = align.align_pose(export_path=data_path)

class AlignPose(pose_data_dict: Dict[str, ndarray], legs_list: List[str], include_claw: bool | None = False, body_template: Dict[str, ndarray] | None = None, body_size: Dict[str, float] | None = None, log_level: Literal['DEBUG', 'INFO', 'WARNING', 'ERROR'] = 'INFO')

Bases: object

Aligns the 3D poses. For the class usage examples, please refer to example_alignment.py

Parameters:

• pose_data_dict (Dict[str, np.ndarray]) –

  3D pose put in a dictionary that has the following structure defined by PTS2ALIGN (see data.py for more details)

  Example format

  >>> pose_data_dict = {
          "RF_leg": np.ndarray[N_frames,N_key_points,3],
          "LF_leg": np.ndarray[N_frames,N_key_points,3],
          "R_head": np.ndarray[N_frames,N_key_points,3],
          "L_head": np.ndarray[N_frames,N_key_points,3],
          "Neck": np.ndarray[N_frames,N_key_points,3],
      }
  

• legs_list (List[str]) – A list containing the leg names to operate on. Should follow the convention <R or L><F or M or H> e.g., [“RF”, “LF”, “RM”, “LM”, “RH”, “LH”]

• include_claw (bool, optional) – If True, claw is included in the scaling process, by default False

• body_template (Dict[str, np.ndarray], optional) – A dictionary containing the positions of fly model body segments. Check ./data.py for the default dictionary, by default None

• body_size (Dict[str, float], optional) – A dictionary containing the limb size of the fly. If the user wants to scale the animal data to match the biomechanical model, then body_size should be the same as the model body size. Otherwise, the user should calculate the animal’s body size. Check ./data.py for an example.

• log_level (Literal["DEBUG", "INFO", "WARNING", "ERROR"], optional) – Logging level as a string, by default “INFO”

align_head(head_array: ndarray, side: str) → ndarray

Scales and translates the head key point locations based on the model size and configuration.

This method takes a 3D array of head key point positions and scales and translates them to align with a predefined model size and configuration, such as a fly”s head. It accounts for the relative positions of key points and ensures that the scaled head key points match the model size.

Parameters:

• head_array (np.ndarray) – A 3D array containing the head key point positions.

• side (str) – A string indicating the side of the head (e.g., “R” or “L”) for alignment.

Returns:

A new 3D array containing the scaled and aligned head key point positions.

Return type:

np.ndarray

Raises:

KeyError – If the “body_template” dictionary does not contain the required key names, “Antenna_mid_thorax” or “Antenna”.

Notes

• This method is used to align head key point positions with a model of a fly”s head.

• It calculates the scale factor and translations necessary to match the model’s head size and position.

align_leg(leg_array: ndarray, leg_name: Literal['RF', 'LF', 'RM', 'LM', 'RH', 'LH']) → ndarray

Scales and translates the leg key point locations based on the model size and configuration.

This method takes a 3D array of leg key point positions and scales and translates them to align with a predefined model size and joint configuration. It accounts for the relative positions of key points and ensures that the scaled leg key points match the model size.

Parameters:

• leg_array (np.ndarray) – A 3D array containing the leg key point positions.

• leg_name (str) – A string indicating the name of the leg (e.g., “RF”, “LF”, …) for alignment.

Returns:

A new 3D array containing the scaled and aligned leg key point positions.

Return type:

np.ndarray

Notes

• This method is used to align leg key point positions with a model of a fly”s leg.

• It calculates the scale factor and multiplies the first 4 or 5 segments with the scale factor.

align_pose(export_path: str | Path | None = None) → Dict[str, ndarray]

Aligns the leg and head key point positions.

Parameters:

export_path (Union[str, Path], optional) – The path where the aligned pose data will be saved, if specified.

Returns:

A dictionary containing the aligned pose data.

Return type:

Dict[str, np.ndarray]

find_scale_leg(leg_name: str, mean_length: Dict) → float

Computes the ratio between the model size and the real fly size.

find_stationary_indices(array: ndarray, threshold: float | None = 5e-05) → ndarray

Find the indices in an array where the function value does not move significantly.

classmethod from_file_path(main_dir: str | Path, file_name: str | None = 'pose3d.*', convert_func: Callable | None = None, pts2align: Dict[str, List[str]] | None = None, **kwargs)

Class method to load pose3d data and convert it into a proper structure.

Parameters:

• main_dir (Union[str, Path]) – Path where the Anipose triangulation results are saved. By default, the result file is caled pose3d.h5 However, if the name is different, <pose_result_path> should be modified accordingly. Example: “../Fly001/001_Beh/behData/pose-3d”

• file_name (str, optional) – File name, by default “pose3d.*”

• convert_func (Callable, optional) – Function to convert the loaded pose into the required format, set by the user if None, no conversion is performed, by default None

• pts2align (Dict[str, List[str]], optional) – Body part names and corresponding key points names to be aligned, check data.py for an example, by default None

Returns:

Instance of the AlignPose class.

Return type:

AlignPose

Raises:

FileNotFoundError – If file with a name that contains file_name does not exist in main_dir.

static get_fixed_pos(points_3d: ndarray) → ndarray

Gets the fixed pose of a steady key point determined by the quantiles.

get_mean_length(segment_array: ndarray, segment_is_leg: bool) → Dict[str, float]

Computes the mean length of a body segment.

property thorax_mid_pts: ndarray

Gets the middle point of right and left wing hinges.

convert_from_anipose_to_dict(pose_3d: Dict[str, ndarray], pts2align: Dict[str, List[str]]) → Dict[str, ndarray]

Loads anipose 3D pose data into a dictionary. See data.py for a mapping from keypoint name to segment name.

Parameters:

• pose_3d (Dict[str, np.ndarray]) –

  3D pose data from anipose. It should have the following format

  >>> pose_3d = {
      "{keypoint_name}_x" : np.ndarray[N_frames,],
      "{keypoint_name}_y" : np.ndarray[N_frames,],
      "{keypoint_name}_z" : np.ndarray[N_frames,],
  }
  

• pts2align (Dict[str, List[str]]) – Segment names and corresponding key point names to be aligned, check data.py for an example, by default None

Returns:

Pose data dictionary of the following format

>>> pose_data_dict = {
    "RF_leg": np.ndarray[N_frames,N_key_points,3],
    "LF_leg": np.ndarray[N_frames,N_key_points,3],
    "R_head": np.ndarray[N_frames,N_key_points,3],
    "L_head": np.ndarray[N_frames,N_key_points,3],
    "Neck": np.ndarray[N_frames,N_key_points,3],
    ...
}

Return type:

Dict[str, np.ndarray]

convert_from_df3d_to_dict(pose_3d: ndarray, pts2align: Dict[str, ndarray]) → Dict[str, ndarray]

Loads DeepFly3D data into a dictionary. See the original DeepFly3D repository for indices of key points for each segment.

Parameters:

• pose_3d (np.ndarray) – Array (N, N_key_points, 3) containing 3D pose data.

• pts2align (Dict[str, np.ndarray]) –

  Dictionary mapping segment names to key point indices. Should be in the following format

  >>> pts2align = {
      "RF_leg": np.arange(0,5),
      "RM_leg": np.arange(5,10),
      "RH_leg": np.arange(10,15),
      "LF_leg": np.arange(19,24),
      "LM_leg": np.arange(24,29),
      "LH_leg": np.arange(29,34),
  }
  

Returns:

Pose data dictionary as described above.

Return type:

Dict[str, np.ndarray]

convert_from_df3dpp_to_dict(pose_3d: Dict[str, Dict[str, ndarray]], pts2align: List[str] | None = None) → Dict[str, ndarray]

Load DeepFly3DPostProcessing data into a dictionary.

Parameters:

• pose_3d (Dict[str, Dict[str, np.ndarray]]) – 3D pose data from DeepFly3DPostProcessing. See the original repository for details.

• pts2align (Optional[List[str]], optional) –

  List of legs to take into account, by default None Example format

  >>> pts2align = ["RF_leg", "LF_leg"]
  

Returns:

Pose data dictionary as described above.

Return type:

Dict[str, np.ndarray]

seqikpy.data module

Data, constants, and paths.

get_pts2align(path: str)

Deletes the keys from the PTS2ALIGN dictionary.

seqikpy.head_inverse_kinematics module

Implementation of a class to calculate head inverse kinematics.

Example usage: >>> import pickle >>> from pathlib import Path >>> from seqikpy.head_inverse_kinematics import HeadInverseKinematics >>> from seqikpy.data import NMF_TEMPLATE

>>> DATA_PATH = Path("../data/anipose/normal_case/pose-3d")
>>> f_path = DATA_PATH / "aligned_pose3d.h5"
>>> with open(f_path, "rb") as f:
>>>     data = pickle.load(f)

>>> class_hk = HeadInverseKinematics(
        aligned_pos = data,
        body_template=NMF_TEMPLATE,
    )
>>> joint_angles = class_hk.compute_head_angles(export_path = DATA_PATH)

IMPORTANT NOTES:

• The aligned pose must include keys: [“R_head”, “L_head”, “Neck”] for head angles calculation.

• Each key corresponds to an array shaped as (N, key_points, 3), where N is the frame count, key_points is the number of key points per part (2 for the head, including antenna base and edge, 1 for the neck), and 3 represents the x, y, z dimensions.

• The key points vary based on the 3D data. If antennal joint angles are not required, “L_head” and “R_head” may contain only a single key point. In such cases, set calculate_ant_angle to False in compute_head_angles. The head segments (“L_head”, “R_head”) can include any key point such as the head bristles for calculating head roll, pitch, and yaw. The examples utilize the antennae base for these calculations.

class AxesTuple(X_AXIS, Y_AXIS, Z_AXIS)

Bases: tuple

X_AXIS

Alias for field number 0

Y_AXIS

Alias for field number 1

Z_AXIS

Alias for field number 2

class HeadInverseKinematics(aligned_pos: Dict[str, ndarray], body_template: Dict[str, ndarray], log_level: Literal['DEBUG', 'INFO', 'WARNING', 'ERROR'] = 'INFO')

Bases: object

Calculates the head DOFs (3) and antennae DOFs (2)

Parameters:

• aligned_pos (Dict[str, np.ndarray]) – Aligned pose dictionary. In principle, it should have body parts (R_head) as keys, and arrays (N,2 for the head) as values. Check the sample data for more detailed example.

• body_template (Dict[str, np.ndarray]) – Dictionary containing the positions of fly model body segments. Check data.py for the default dictionary.

• log_level (Literal["DEBUG", "INFO", "WARNING", "ERROR"], optional) – Logging level as a string, by default “INFO”

static angle_between_segments(v1: ndarray, v2: ndarray, rot_axis: ndarray) → float

Calculates the angle between two vectors based on the cos formula. It reverses the sign of the angle if determinant of the matrix having two vectors and the rotation axis is negative.

The returned angle is in radians.

compute_antenna_pitch(side: Literal['R', 'L'], head_roll: ndarray) → ndarray

Calculates the head pitch angle (rad) from head vector projected onto sagittal plane to antenna vector (from base ant to edge). Furthermore, it subtracts the angle with the resting joint angle of the antenna pitch.

Higher antenna pitch means antenna is lifted upward more.

compute_antenna_yaw(side: Literal['R', 'L'], head_roll: ndarray) → ndarray

Calculates the antenna yaw angle (rad) from the lateral head vector projected onto transverse plane to antenna vector (from base ant to edge) projected, again, on the transverse plane.

Higher antenna yaw means antennae get closer to the midline, deviating from the resting position.

compute_head_angles(export_path: str | Path | None = None, compute_ant_angles: bool | None = True) → Dict[str, ndarray]

Calculates the head & antennal joint angles

Parameters:

• export_path (Union[str, Path], optional) – Export path of the joint angles, by default None

• compute_ant_angles (Optional[bool], optional) – If True, computes the head roll,pitch and yaw, by default True

Returns:

Dicitonary containing the head joint angles, saved in export path if provided

Return type:

Dict[str, np.ndarray]

compute_head_pitch() → ndarray

Calculates the head pitch angle (rad) from head mid vector projected onto sagittal plane to the anteroposterior plane. Furthermore, it sums the angle with the resting joint angle of the head pitch.

Higher head pitch means head is lowered more.

compute_head_roll() → ndarray

Calculates the head roll angle (rad) from horizontal axis to head horizontal vector projected onto transverse plane.

Positive head roll -> rotation to the right in fly coords Negative head roll -> rotation to the left in the fly coords

compute_head_yaw() → ndarray

Calculates the head yaw angle (rad) from horizontal axis to head horizontal vector projected onto frontal plane.

Positive head yaw -> head yaw to the left Negative head yaw -> head yaw to the right

derotate_vector(head_roll_angle: float, vector_to_derotate: ndarray) → ndarray

Rotates a vector by the inverse amount of head_roll_angle along the x axis.

get_ant_vector(side: Literal['R', 'L']) → ndarray

Vector ((N,3) array) from antenna base to antenna edge.

get_head_vector(side: Literal['R', 'L']) → ndarray

Vector ((N,3) array) from <side> antenna base (or any head key point) to neck.

get_head_vector_horizontal() → ndarray

Vector ((N,3) array) from right antenna base (or any head key point) to left antenna base (or any head key point).

get_head_vector_mid() → ndarray

Vector ((N,3) array) from mid antenna base (or any head key point) to neck.

get_plane(row: ndarray, n_row: int) → ndarray

Construct an array by repeating row n_row many times.

get_rest_antenna_pitch() → float

Antenna pitch angle at zero pose in the fly biomechanical model.

get_rest_head_pitch() → float

Head pitch angle at zero pose in the fly biomechanical model.

seqikpy.kinematic_chain module

Module that contains a set of kinematic chains.

class AxesTuple(X_AXIS, Y_AXIS, Z_AXIS)

Bases: tuple

X_AXIS

Alias for field number 0

Y_AXIS

Alias for field number 1

Z_AXIS

Alias for field number 2

class KinematicChainBase(bounds_dof: Dict[str, ndarray], legs_list: List[str], body_size: Dict[str, float] | None = None)

Bases: ABC

Abstract class to create kinematic chains for the legs.

Parameters:

• bounds_dof (Dict[str, np.ndarray]) – Dictionary that contains the bounds of joint degrees-of-freedom.

• legs_list (List[str]) – List of legs for which the kinematic chains are created.

• body_size (Dict[str, float], optional) – Dictionary that contains the size of different body parts, by default None

abstract create_leg_chain(leg_name: str, **kwargs) → Chain

Returns the respective leg chain based on the stage.

Parameters:

leg_name (str) – Name of the leg, RF or LF

Return type:

KinematicChain of a given leg

class KinematicChainGeneric(bounds_dof: Dict[str, ndarray], legs_list: List[str], body_size: Dict[str, float] | None = None)

Bases: KinematicChainBase

Generic kinematic chain class to create one kinematic chain for the entire leg.

Parameters:

• bounds_dof (Dict[str, np.ndarray]) – Dictionary that contains the bounds of joint degrees-of-freedom.

• legs_list (List[str]) – List of legs for which the kinematic chains are created.

• body_size (Dict[str, float], optional) – Dictionary that contains the size of different body parts, by default None

create_leg_chain(leg_name: str, **kwargs) → Chain

Returns the respective leg chain based on the stage.

Parameters:

• leg_name (str) – Name of the leg, RF or LF

• angles (Dict[str, np.ndarray], optional) – Joint angles calculated in the previous step, None at the first step, by default None

Return type:

The entire KinematicChain of a given leg

class KinematicChainSeq(bounds_dof: Dict[str, ndarray], legs_list: List[str], body_size: Dict[str, float] | None = None)

Bases: KinematicChainBase

Sequential kinematic chain class to create a kinematic chain stage-by-stage.

Parameters:

• bounds_dof (Dict[str, np.ndarray]) – Dictionary that contains the bounds of joint degrees-of-freedom.

• legs_list (List[str]) – List of legs for which the kinematic chains are created.

• body_size (Dict[str, float], optional) – Dictionary that contains the size of different body parts, by default None

Note

This implementation follows Yaw-Pitch-Roll order.

create_leg_chain(leg_name: str, **kwargs) → Chain

Returns the respective leg chain based on the stage.

Parameters:

• leg_name (str) – Name of the leg, i.e., RF, LF, etc.

• (kwargs) (t) – Joint angles calculated in the previous step, None at the first step, by default None

• (kwargs) – Stage number, (1,2,3,4)

• (kwargs) – Time step, by default 0

Return type:

KinematicChain at the respective stage and time step.

Raises:

• ValueError – If the stage number is not in (1,2,3,4)

• ValueError – If the leg name is not in [“RF”, “LF”, “RM”, “LM”, “RH”, “LH”]

create_leg_chain_stage_1(leg_name: str) → Chain

Leg chain to calculate thorax/coxa pitch and yaw.

Parameters:

leg_name (str) – Leg name, e.g., RF, LF, RM, LM, RH, LH

Returns:

Chain of the leg at stage 1, which calculates thorax/coxa pitch and yaw joint angles and only contains Coxa leg segment.

Return type:

Chain

create_leg_chain_stage_2(leg_name: str, angles: Dict[str, ndarray], t: int) → Chain

Leg chain to calculate thorax/coxa roll and coxa/femur pitch.

Parameters:

leg_name (str) – Leg name, e.g., RF, LF, RM, LM, RH, LH

Returns:

Chain of the leg at stage 2, which calculates thorax/coxa roll and coxa/femur pitch joint angles and contains Coxa and Femur leg segments.

Return type:

Chain

create_leg_chain_stage_3(leg_name: str, angles: Dict[str, ndarray], t: int) → Chain

Leg chain to calculate coxa/femur roll and femur/tibia pitch.

Parameters:

leg_name (str) – Leg name, e.g., RF, LF, RM, LM, RH, LH

Returns:

Chain of the leg at stage 3, which calculates coxa/femur roll and femur/tibia pitch joint angles and contains Coxa, Femur and Tibia leg segments.

Return type:

Chain

create_leg_chain_stage_4(leg_name: str, angles: Dict[str, ndarray], t: int) → Chain

Leg chain to calculate tibia/tarsus pitch.

Parameters:

leg_name (str) – Leg name, e.g., RF, LF, RM, LM, RH, LH

Returns:

Chain of the leg at stage 4, which calculates coxa/femur roll and tibia/tarsus pitch joint angles and contains the entire leg, i.e., Coxa, Femur, Tibia and Tarsus.

Return type:

Chain

seqikpy.leg_inverse_kinematics module

Module to calculate inverse kinematics from 3D pose based on IKPy.

class LegInvKinBase(aligned_pos: Dict[str, ndarray], kinematic_chain_class: KinematicChainBase, initial_angles: Dict[str, ndarray] | None = None, log_level: Literal['DEBUG', 'INFO', 'WARNING', 'ERROR'] = 'INFO')

Bases: ABC

Abstract class to calculate inverse kinematics for leg joints.

Parameters:

• aligned_pos (Dict[str, np.ndarray]) –

  Aligned pose from the AlignPose class. Should have the following structure

  >>> pose_data_dict = {
          "<side><segment>_leg": np.ndarray[N_frames,N_key_points,3],
      }
  

• kinematic_chain (KinematicChainBase) – Kinematic chain of the leg.

• initial_angles (Dict[str, np.ndarray], optional) – Initial angles of DOFs. If not provided, the default values from data.py will be used.

• log_level (Literal["DEBUG", "INFO", "WARNING", "ERROR"], optional) – Logging level as a string, by default “INFO”

calculate_fk(kinematic_chain: Chain, joint_angles: ndarray) → ndarray

Calculates the forward kinematics from the joint dof angles.

calculate_ik(kinematic_chain: Chain, target_pos: ndarray, initial_angles: ndarray | None = None) → ndarray

Calculates the joint angles in the leg chain.

abstract calculate_ik_stage(end_effector_pos: ndarray, origin: ndarray, initial_angles: ndarray, segment_name: str, **kwargs) → ndarray

For a given trial pose data, calculates the inverse kinematics.

Parameters:

• end_effector_pos (np.ndarray) – 3D array containing the position of the end effector pos

• origin (np.ndarray) – Origin of the kinematic chain, i.e., Thorax-Coxa joint

• initial_angles (np.ndarray) – Initial angles for the optimization seed

• segment_name (str) – Leg side, i.e., RF, LF, RM, LM, RH, LH

Returns:

Array containing the cartesian coordinates of the joint positions. The joint angles are saved in a class attribute

Return type:

np.ndarray

get_scale_factor(vector: ndarray, length: float) → float

Gets scale the ratio between two vectors.

abstract run_ik_and_fk(export_path: str | Path | None = None, **kwargs) → Tuple[Dict[str, ndarray], Dict[str, ndarray]]

Runs inverse and forward kinematics for leg joints.

Parameters:

export_path (Union[Path, str], optional) – Path where the results will be saved, if None, nothing is saveed, by default None

Returns:

Two dictionaries containing joint angles and forward kinematics, respectively.

Return type:

Tuple[Dict[str,np.ndarray], Dict[str,np.ndarray]]

class LegInvKinGeneric(aligned_pos: Dict[str, ndarray], kinematic_chain_class: KinematicChainGeneric, initial_angles: Dict[str, ndarray] | None = None, log_level: Literal['DEBUG', 'INFO', 'WARNING', 'ERROR'] = 'INFO')

Bases: LegInvKinBase

Class to calculate inverse kinematics for leg joints in a generic manner to only match the given end effector.

Parameters:

• aligned_pos (Dict[str, np.ndarray]) –

  Aligned pose from the AlignPose class. Should have the following structure

  >>> pose_data_dict = {
          "<side><segment>_leg": np.ndarray[N_frames,N_key_points,3],
      }
  

• kinematic_chain (KinematicChainGeneric) – Kinematic chain of the leg.

• initial_angles (Dict[str, np.ndarray], optional) – Initial angles of DOFs. If not provided, the default values from data.py will be used.

• log_level (Literal["DEBUG", "INFO", "WARNING", "ERROR"], optional) – Logging level as a string, by default “INFO”

Examples

>>> from pathlib import Path
>>> from seqikpy.kinematic_chain import KinematicChainGeneric
>>> from seqikpy.leg_inverse_kinematics import LegInvKinGeneric
>>> from seqikpy.data import BOUNDS, INITIAL_ANGLES
>>> from seqikpy.utils import load_file

>>> DATA_PATH = Path("../data/anipose_220525_aJO_Fly001_001/pose-3d")
>>> f_path = DATA_PATH / "pose3d_aligned.pkl"

>>> aligned_pos = load_file(f_path)

>>> seq_ik = LegInvKinGeneric(
        aligned_pos=aligned_pos,
        kinematic_chain_class=KinematicChainGeneric(
            bounds_dof=BOUNDS,
            legs_list=["RF", "LF"],
            body_size=None,
        ),
        initial_angles=INITIAL_ANGLES
    )

>>> leg_joint_angles, forward_kinematics = gen_ik.run_ik_and_fk(
        export_path=DATA_PATH,
        hide_progress_bar=False
    )

calculate_ik_stage(end_effector_pos: ndarray, origin: ndarray, initial_angles: ndarray, segment_name: str, **kwargs) → ndarray

For a given trial pose data, calculates the inverse kinematics.

Parameters:

• end_effector_pos (np.ndarray) – 3D array containing the position of the end effector pos

• origin (np.ndarray) – Origin of the kinematic chain, i.e., Thorax-Coxa joint

• initial_angles (np.ndarray) – Initial angles for the optimization seed

• segment_name (str) – Leg side, i.e., RF, LF, RM, LM, RH, LH

Returns:

Array containing the cartesian coordinates of the joint positions. The joint angles are saved in a class attribute

Return type:

np.ndarray

run_ik_and_fk(export_path: str | Path | None = None, **kwargs) → Tuple[Dict[str, ndarray], Dict[str, ndarray]]

Runs inverse and forward kinematics for leg joints.

Parameters:

export_path (Union[Path, str], optional) – Path where the results will be saved, if None, nothing is saveed, by default None

Returns:

Two dictionaries containing joint angles and forward kinematics, respectively.

Return type:

Tuple[Dict[str,np.ndarray], Dict[str,np.ndarray]]

class LegInvKinSeq(aligned_pos: Dict[str, ndarray], kinematic_chain_class: KinematicChainSeq, initial_angles: Dict[str, ndarray] | None = None, log_level: Literal['DEBUG', 'INFO', 'WARNING', 'ERROR'] = 'INFO')

Bases: LegInvKinBase

Class to calculate inverse kinematics for leg joints in a sequential manner. This method finds the optimal joint angles within the bounds to match each koint as closely as possible.

Parameters:

• aligned_pos (Dict[str, np.ndarray]) –

  Aligned pose from the AlignPose class. Should have the following structure

  >>> pose_data_dict = {
          "<side><segment>_leg": np.ndarray[N_frames,N_key_points,3],
      }
  

• kinematic_chain (KinematicChainSeq) – Kinematic chain of the leg.

• initial_angles (Dict[str, np.ndarray], optional) – Initial angles of DOFs. If not provided, the default values from data.py will be used.

• log_level (Literal["DEBUG", "INFO", "WARNING", "ERROR"], optional) – Logging level as a string, by default “INFO”

Examples

>>> from pathlib import Path
>>> from seqikpy.kinematic_chain import KinematicChainSeq
>>> from seqikpy.leg_inverse_kinematics import LegInvKinSeq
>>> from seqikpy.data import BOUNDS, INITIAL_ANGLES
>>> from seqikpy.utils import load_file

>>> DATA_PATH = Path("../data/anipose_220525_aJO_Fly001_001/pose-3d")
>>> f_path = DATA_PATH / "pose3d_aligned.pkl"

>>> aligned_pos = load_file(f_path)

>>> seq_ik = LegInvKinSeq(
        aligned_pos=aligned_pos,
        kinematic_chain_class=KinematicChainSeq(
            bounds_dof=BOUNDS,
            legs_list=["RF", "LF"],
            body_size=None,
        ),
        initial_angles=INITIAL_ANGLES
    )

>>> leg_joint_angles, forward_kinematics = seq_ik.run_ik_and_fk(
        export_path=DATA_PATH,
        hide_progress_bar=False
    )

calculate_ik_stage(end_effector_pos: ndarray, origin: ndarray, initial_angles: ndarray, segment_name: str, **kwargs) → ndarray

For a given trial pose data, calculates the inverse kinematics.

Parameters:

• end_effector_pos (np.ndarray) – 3D array containing the position of the end effector pos

• origin (np.ndarray) – Origin of the kinematic chain, i.e., Thorax-Coxa joint

• initial_angles (np.ndarray) – Initial angles for the optimization seed

• (kwargs) (stage) – Leg side, i.e., RF, LF, RM, LM, RH, LH

• (kwargs) – Stage in the sequential calculation, should be between 1 and 4

Returns:

Array containing the cartesian coordinates of the joint positions. The joint angles are saved in a class attribute

Return type:

np.ndarray

run_ik_and_fk(export_path: str | Path | None = None, **kwargs) → Tuple[Dict[str, ndarray], Dict[str, ndarray]]

Runs inverse and forward kinematics for leg joints.

Parameters:

• export_path (Union[Path, str], optional) – Path where the results will be saved, if None, nothing is saveed, by default None

• (kwargs) (hide_progress_bar) – Stages to run the inverse kinematics.

• (kwargs) – Hide the progress bar, by default True

Returns:

Two dictionaries containing joint angles and forward kinematics, respectively.

Return type:

Tuple[Dict[str,np.ndarray], Dict[str,np.ndarray]]

seqikpy.utils module

Utilities.

calculate_body_size(body_template: Dict[str, ndarray], legs_list: List[str] = ['RF', 'LF', 'RM', 'LM', 'RH', 'LH']) → Dict[str, ndarray]

Calculates body segment sizes from the template data.

compute_length_of_segment(points3d, segment_beg, segment_end)

Computes the length of a segment.

compute_mean_length_of_segments(segment_lengths)

Computes mean length of each segment.

df_to_nparray(data_frame, side, claw_is_end_effector, segment='F')

Convert usual dataframe format into a three dimensional array.

dict_to_nparray_angle(angle_dict, leg, claw_is_end_effector)

Convert usual df3dPP dictionary format into a three dimensional array.

dict_to_nparray_pose(pose_dict, claw_is_end_effector)

Convert usual df3dPP dictionary format into a three dimensional array.

dist_calc(v1, v2)

drop_level_dlc(data_frame)

Converts DLC type dataframe into one level df.

fix_coxae_pos(points3d, right_coxa_kp='thorax_coxa_R', left_coxa_kp='thorax_coxa_L')

Calculates the fixed coxae location based on the quantiles.

from_anipose_to_array(points3d, claw_is_end_effector=False, kps=['thorax_coxa', 'coxa_femur', 'femur_tibia', 'tibia_tarsus']) → ndarray

Convert usual dataframe format into a three dimensional array of size (N,KeyPoints,3).

from_sdf(sdf_file: str)

Extracts a body template and joint bounds from an SDF file.

Parameters:

sdf_file (str) – Path to the SDF file.

Returns:

Body template and joint bounds.

Return type:

Dict, Dict

get_array(kp_name, kp_dict)

Returns 3D points of a key point in an array format.

get_distance_btw_vecs(vector1, vector2)

get_fps_from_video(video_dir)

Finds the fps of a video.

get_leg_length(mean_segment_size)

Returns the leg size from the mean segment lengths.

get_length_of_segments(points3d, claw_is_ee=False)

Returns a dictionary with segment sizes.

get_mean_length_of_segments(points3d)

Gets a dictionary with segment lengths.

get_mean_quantile(vector, quantile_diff=0.05)

get_stim_array(lines: list, frame_rate: int, scale: int = 1, hamming_window_size: int = 0, time_scale: int = 0.001)

Computes a stimulus array from the stimulus information.

Parameters:

• main_dir (str) – Path of the file containing the stim info.

• frame_rate (int) – Frame rate of the recordings.

• scale (int, optional) – Interpolation rate applied in df3dPP, by default 1 (no interpolation!)

• hamming_window_size (int, optional) – Size of the filter applied in df3dPP, by default 28

• time_scale (int, optional) – Time scale of the stimulus info (msec), by default 1e-3

Returns:

np.array(length, ) of boolean values indicating the stimulus time points.

Return type:

np.ndarray

get_stim_intervals(stim_data)

Reads stimulus array and returns the stim intervals for plotting purposes. Use get_stim_array otherwise.

interpolate_joint_angles(joint_angles_dict, **kwargs)

Interpolates joint angles.

interpolate_signal(signal, original_ts, new_ts)

Interpolates signals.

leg_length_model(nmf_size: dict, leg_name: str, claw_is_ee: bool) → float

Returns the leg size from the nmf size dictionary.

load_file(output_fname)

Load file.

load_stim_data(main_dir: Path)

Loads the stimulus info from txt.

save_file(out_fname, data)

Save file.

seqikpy.visualization module

Functions for plotting and animation.

animate_3d_points(points3d: Dict[str, ndarray], export_path: Path, points3d_second: Dict[str, ndarray] | None = None, fps: int = 100, frame_no: int = 1000, format_video: str = 'mp4', elev: int = 10, azim: int = 90, title: str = '', marker_types: Dict[str, str] | None = None) → None

Makes an animation of 3D pose. This code is intended for animating the raw 3D pose and forward kinematics from seqikpy.

Parameters:

• points3d (Dict[str, np.ndarray]) – Dictionary containing the 3D pose, usually this is the raw 3D pose.

• export_path (Path) – Path where the animation will be saved.

• points3d_second (Dict[str, np.ndarray], optional) – Dictionary containing the 3D pose usually this is the forward kinematics, by default None

• fps (int, optional) – Frames per second, by default 100

• frame_no (int, optional) – Maximum number of frames of the animation, by default 1000

• format_video (str, optional) – Video format, by default “mp4”

• elev (int, optional) – Elevation of the point of view, by default 10

• azim (int, optional) – Azimuth of the point of view, by default 90

• title (str, optional) – Title of the video, by default “”

• marker_types (Dict[str, str], optional) – Marker types for each key point, by default None

fig_to_array(fig)

Converts a matplotlib figure into an array.

generate_color_map(cmap: str, n: int) → List

Generates a list of colors from a given colormap.

get_frames_from_video_ffmpeg(path)

Saves frames of a video using FFMPEG.

Parameters:

path (Path) – Video path. This path appended with a _frames folder will be used to save the frames.

get_plot_config(data_path: Path)

Get experimental conditions from the data path. Data path should look like: “/mnt/nas2/GO/7cam/220810_aJO-GAL4xUAS-CsChr/Fly001/001_RLF/behData/pose-3d”

get_video_writer(video_path, fps, output_shape)

Initialize and return a VideoWriter object.

load_grid_plot_data(data_path: Path) → [Dict, Dict]

Loads the set of data necessary for plotting the grid.

Parameters:

data_path (Path) – Data path where the pose3d and inverse kinematics are saved.

Returns:

Returns joint angles (head and leg) and aligned pose as a tuple.

Return type:

Tuple

make_video(video_path: str, frame_generator: Iterable, fps: int, output_shape: Tuple = (-1, 2880), n_frames: int = -1)

Makes videos from a generator of images.

plot_3d_points(ax3d, points3d, export_path=None, t=0, marker_types=None, line_style='solid')

Plots 3D points at time t.

plot_grid(img_front: ndarray, img_side: ndarray, aligned_pose: Dict[str, ndarray], joint_angles: Dict[str, ndarray], leg_angles_to_plot: List[str], head_angles_to_plot: List[str], t: int, t_start: int, t_end: int, t_interval: int = 20, fps: int = 100, trail: int = 30, key_points_to_trail: Dict[str, Iterable] | None = None, marker_types_3d: Dict[str, str] | None = None, marker_trail: str | None = None, stim_lines: List[int] | None = None, export_path: Path | None = None, **kwargs)

Plots an instance of the animal recording, 3D pose, head and leg joint angles in a grid layout. This code is intended to use in a for loop to plot and save all the frames from t_start to t_end to make a video afterwards.

NOTE: This function is intended to plot leg and head joint angles together, it will not work if any of these data is missing.

Parameters:

• img_front (Path) – Image of the fly at frame t on the front camera.

• img_side (Path) – Image of the fly at frame t on the side camera.

• aligned_pose (Dict[str, np.ndarray]) – Aligned 3D pose.

• joint_angles (Dict[str, np.ndarray]) – Joint angles.

• leg_angles_to_plot (List[str]) –

  List containing leg joint angle names without the side. Examples

  >>> leg_joint_angles = [
          "ThC_yaw",
          "ThC_pitch",
          "ThC_roll",
          "CTr_pitch",
          "CTr_roll",
          "FTi_pitch",
          "TiTa_pitch",
          ]
  

• head_angles_to_plot (List[str]) – List containing exact names of head joint angle names.

• t (int) – Frame number t.

• t_start (int) – Start of the time series, i.e., joint angles.

• t_end (int) – End of the time series, i.e., joint angles.

• t_interval (int, optional) – Interval of frame numbers in between x ticks, by default 20

• fps (int, optional) – Frames per second, by default 100

• trail (int, optional) – Number of previous frames where the key point will be visible, by default 30

• marker_types_3d (Dict[str, Tuple[np.ndarray, str]]) – Dictionary mapping key points names to marker styles.

• marker_trail (Dict[str, Tuple[np.ndarray, str]]) – Dictionary mapping key points names to their indices and line styles for trailing key points.

• stim_lines (List[int], optional) – Stimulation indicators, by default None

• export_path (Path, optional) – Path where the plot will be saved, by default None

Returns:

Figure.

Return type:

Fig

plot_grid_generator(fly_frames_front: ndarray, fly_frames_side: ndarray, aligned_pose: Dict[str, ndarray], joint_angles: Dict[str, ndarray], leg_angles_to_plot: List[str], head_angles_to_plot: List[str], t_start: int, t_end: int, t_interval: int = 20, fps: int = 100, trail: int = 30, key_points_to_trail: List[str] | None = None, marker_types_3d: Dict[str, str] | None = None, marker_trail: str | None = 'x', stim_lines: List[int] | None = None, export_path: Path | None = None, **kwargs)

Generator for plotting grid.

plot_joint_angle(ax: Axes, kinematics_data: Dict[str, ndarray], angles_to_plot: List, degrees: bool = True, until_t: int = -1, stim_lines: List[int] | None = None, show_legend: bool = True, export_path: Path | None = None)

Plot joint angles from a given kinematics data.

Parameters:

• ax (plt.Axes) – Axis where the plot will be displayed.

• kinematics_data (Dict[str, np.ndarray]) – Dictionary containing the kienmatics, pose and angle

• angles_to_plot (List) – Angles to plot. Exact column name should be given.

• degrees (bool, optional) – Convert to degrees, by default True

• until_t (int, optional) – Plots the angles until t^th frame, by default -1

• stim_lines (List[int], optional) – Plots vertical lines in indicated locations, by default None

• show_legend (bool, optional) – Shows legend, by default True

• export_path (Path, optional) – Path where the plot will be saved, by default None

plot_trailing_kp(ax3d, points3d, segments_to_plot, export_path=None, t=0, trail=5, marker_type='x')

Plots the traces of key points from t-trail to t.

resize_frame(original_size, target_size)

Resize a frame to the target size while preserving the width length ratio.

resize_rgb_frame(frame, output_shape)

resize the frame and convert it to RGB

video_frames_generator(video_path: Path, start_frame: int, end_frame: int, stim_lines: List[int], radius=30, center=(50, 50), color=(255, 0, 0))

Returns the frames as a generator in a given interval. Modifies the brightness and contrast of the images.

Parameters:

• video_path (Path) – Video path.

• start_frame (int) – Starting frame.

• end_frame (int) – End frame.

Yields:

Frame – Generator containing all the frames in the specified interval.

Module contents