Anatomy flybody

FlyBody-specific anatomical definitions.

This module extends the default anatomy types from flygym.anatomy.

FlyBodyActuatedDOFPreset

Bases: BaseActuatedDOFPreset

Presets for which flybody joint DoFs should be actuated.

Source code in src/flygym/flybody/anatomy_flybody.py

class FlyBodyActuatedDOFPreset(BaseActuatedDOFPreset):
    """Presets for which flybody joint DoFs should be actuated."""

    ALL = "all"
    LEGS_ONLY = "legs_only"
    LEGS_ACTIVE_ONLY = "legs_active_only"

    @classmethod
    def _get_passive_tarsal_links(cls) -> list[str]:
        return FLYBODY_PASSIVE_TARSAL_LINKS

FlyBodyAnatomicalJoint dataclass

Bases: AnatomicalJoint

Anatomical joint specific to the flybody model.

Source code in src/flygym/flybody/anatomy_flybody.py

class FlyBodyAnatomicalJoint(AnatomicalJoint):
    """Anatomical joint specific to the flybody model."""

    def iter_dofs(self, axis_order: AxisOrder) -> Iterator[FlyBodyJointDOF]:
        """Iterate through the DOFs of this joint in the specified axis order."""
        if self.child.is_wing():
            wing_axis_order = WingFlyBodyAxisOrder(
                [axis.value for axis in axis_order.value]
            )
            for axis in wing_axis_order.value:
                if axis in self.axes:
                    yield FlyBodyJointDOF(
                        parent=self.parent,
                        child=self.child,
                        axis=axis,
                    )
            return

        for axis in axis_order.value:
            if axis in self.axes:
                yield FlyBodyJointDOF(
                    parent=self.parent,
                    child=self.child,
                    axis=axis,
                )

iter_dofs(axis_order)

Iterate through the DOFs of this joint in the specified axis order.

Source code in src/flygym/flybody/anatomy_flybody.py

def iter_dofs(self, axis_order: AxisOrder) -> Iterator[FlyBodyJointDOF]:
    """Iterate through the DOFs of this joint in the specified axis order."""
    if self.child.is_wing():
        wing_axis_order = WingFlyBodyAxisOrder(
            [axis.value for axis in axis_order.value]
        )
        for axis in wing_axis_order.value:
            if axis in self.axes:
                yield FlyBodyJointDOF(
                    parent=self.parent,
                    child=self.child,
                    axis=axis,
                )
        return

    for axis in axis_order.value:
        if axis in self.axes:
            yield FlyBodyJointDOF(
                parent=self.parent,
                child=self.child,
                axis=axis,
            )

FlyBodyAxesSet

Bases: AxesSet

Set of rotation axes using FlyBody's axis convention.

Source code in src/flygym/flybody/anatomy_flybody.py

class FlyBodyAxesSet(AxesSet):
    """Set of rotation axes using FlyBody's axis convention."""

    rotation_axis_class = FlyBodyRotationAxis

FlyBodyAxisOrder

Bases: BaseAxisOrder, Enum

Axis order enum based on FlyBodyRotationAxis.

Source code in src/flygym/flybody/anatomy_flybody.py

class FlyBodyAxisOrder(BaseAxisOrder, Enum):
    """Axis order enum based on FlyBodyRotationAxis."""

    @classmethod
    def _axis_enum_cls(cls):
        return FlyBodyRotationAxis

    PITCH_ROLL_YAW = (
        FlyBodyRotationAxis.PITCH,
        FlyBodyRotationAxis.ROLL,
        FlyBodyRotationAxis.YAW,
    )
    PRY = PITCH_ROLL_YAW
    PITCH_YAW_ROLL = (
        FlyBodyRotationAxis.PITCH,
        FlyBodyRotationAxis.YAW,
        FlyBodyRotationAxis.ROLL,
    )
    PYR = PITCH_YAW_ROLL
    ROLL_PITCH_YAW = (
        FlyBodyRotationAxis.ROLL,
        FlyBodyRotationAxis.PITCH,
        FlyBodyRotationAxis.YAW,
    )
    RPY = ROLL_PITCH_YAW
    ROLL_YAW_PITCH = (
        FlyBodyRotationAxis.ROLL,
        FlyBodyRotationAxis.YAW,
        FlyBodyRotationAxis.PITCH,
    )
    RYP = ROLL_YAW_PITCH
    YAW_PITCH_ROLL = (
        FlyBodyRotationAxis.YAW,
        FlyBodyRotationAxis.PITCH,
        FlyBodyRotationAxis.ROLL,
    )
    YPR = YAW_PITCH_ROLL
    YAW_ROLL_PITCH = (
        FlyBodyRotationAxis.YAW,
        FlyBodyRotationAxis.ROLL,
        FlyBodyRotationAxis.PITCH,
    )
    YRP = YAW_ROLL_PITCH

    DONTCARE = PITCH_ROLL_YAW

FlyBodyBodySegment dataclass

Bases: BodySegment

FlyBody-specific body segment class.

Source code in src/flygym/flybody/anatomy_flybody.py

class FlyBodyBodySegment(BodySegment):
    """FlyBody-specific body segment class."""

    def __post_init__(self):
        if self.name not in FLYBODY_ALL_SEGMENT_NAMES:
            raise ValueError(
                f"Invalid body segment name: {self.name}. "
                f"Must be one of {FLYBODY_ALL_SEGMENT_NAMES}."
            )

    def is_proboscis(self) -> bool:
        """Return True if this segment belongs to the proboscis."""
        return self.link in FLYBODY_PROBOSCIS_LINKS

    def is_eye(self) -> bool:
        """No eyes in flybody model, eyes are part of the head."""
        return False

    def is_antenna(self) -> bool:
        """Return True if this segment belongs to an antenna."""
        return self.link == "antenna"

    def is_leg(self) -> bool:
        """Return True if this segment belongs to a leg."""
        return self.pos in LEGS

    def is_abdomen(self) -> bool:
        """Return True if this segment belongs to the abdomen."""
        return self.link in FLYBODY_ABDOMEN_LINKS

is_abdomen()

Return True if this segment belongs to the abdomen.

Source code in src/flygym/flybody/anatomy_flybody.py

def is_abdomen(self) -> bool:
    """Return True if this segment belongs to the abdomen."""
    return self.link in FLYBODY_ABDOMEN_LINKS

is_antenna()

Return True if this segment belongs to an antenna.

Source code in src/flygym/flybody/anatomy_flybody.py

def is_antenna(self) -> bool:
    """Return True if this segment belongs to an antenna."""
    return self.link == "antenna"

is_eye()

No eyes in flybody model, eyes are part of the head.

Source code in src/flygym/flybody/anatomy_flybody.py

def is_eye(self) -> bool:
    """No eyes in flybody model, eyes are part of the head."""
    return False

is_leg()

Return True if this segment belongs to a leg.

Source code in src/flygym/flybody/anatomy_flybody.py

def is_leg(self) -> bool:
    """Return True if this segment belongs to a leg."""
    return self.pos in LEGS

is_proboscis()

Return True if this segment belongs to the proboscis.

Source code in src/flygym/flybody/anatomy_flybody.py

def is_proboscis(self) -> bool:
    """Return True if this segment belongs to the proboscis."""
    return self.link in FLYBODY_PROBOSCIS_LINKS

FlyBodyContactBodiesPreset

Bases: BaseContactBodiesPreset

Presets for which flybody segments can collide with the ground.

Source code in src/flygym/flybody/anatomy_flybody.py

class FlyBodyContactBodiesPreset(BaseContactBodiesPreset):
    """Presets for which flybody segments can collide with the ground."""

    ALL = "all"
    LEGS_THORAX_ABDOMEN_HEAD = "legs_thorax_abdomen_head"
    LEGS_ONLY = "legs_only"
    TIBIA_TARSUS_ONLY = "tibia_tarsus_only"

    @classmethod
    def _get_all_segments(cls):
        return [FlyBodyBodySegment(segname) for segname in FLYBODY_ALL_SEGMENT_NAMES]

FlyBodyJointDOF dataclass

Bases: JointDOF

Joint DOF specific to the flybody model.

Source code in src/flygym/flybody/anatomy_flybody.py

class FlyBodyJointDOF(JointDOF):
    """Joint DOF specific to the flybody model."""

    @classmethod
    def from_name(cls, name: str) -> "FlyBodyJointDOF":
        """Create a FlyBodyJointDOF from a name of the form 'parent-child-axis'."""
        try:
            parent, child, axis = name.split("-")
            # check if child is wing
            bs_child = FlyBodyBodySegment(child)
            return cls(
                parent=FlyBodyBodySegment(parent),
                child=bs_child,
                axis=FlyBodyRotationAxis(axis)
                if not bs_child.is_wing()
                else WingFlyBodyRotationAxis(axis),
            )
        except ValueError:
            raise ValueError(f"Invalid joint DOF name: {name}. ")

from_name(name) classmethod

Create a FlyBodyJointDOF from a name of the form 'parent-child-axis'.

Source code in src/flygym/flybody/anatomy_flybody.py

@classmethod
def from_name(cls, name: str) -> "FlyBodyJointDOF":
    """Create a FlyBodyJointDOF from a name of the form 'parent-child-axis'."""
    try:
        parent, child, axis = name.split("-")
        # check if child is wing
        bs_child = FlyBodyBodySegment(child)
        return cls(
            parent=FlyBodyBodySegment(parent),
            child=bs_child,
            axis=FlyBodyRotationAxis(axis)
            if not bs_child.is_wing()
            else WingFlyBodyRotationAxis(axis),
        )
    except ValueError:
        raise ValueError(f"Invalid joint DOF name: {name}. ")

FlyBodyRotationAxis

Bases: BaseRotationAxis

FlyBody axis convention.

yaw -> z, pitch -> x, roll -> y.

Source code in src/flygym/flybody/anatomy_flybody.py

class FlyBodyRotationAxis(BaseRotationAxis):
    """FlyBody axis convention.

    yaw -> z, pitch -> x, roll -> y.
    """

    PITCH = "pitch"
    P = PITCH
    ROLL = "roll"
    R = ROLL
    YAW = "yaw"
    Y = YAW

    @classmethod
    def _vector_by_axis(cls) -> dict[str, tuple[float, float, float]]:
        return {
            "pitch": (1, 0, 0),
            "roll": (0, 1, 0),
            "yaw": (0, 0, 1),
        }

FlyBodySkeleton

Bases: Skeleton

Skeleton specific to the flybody model.

Source code in src/flygym/flybody/anatomy_flybody.py

class FlyBodySkeleton(Skeleton):
    """Skeleton specific to the flybody model."""

    def __init__(
        self,
        *,
        axis_order: FlyBodyAxisOrder
        | WingFlyBodyAxisOrder
        | AxisOrder
        | list[RotationAxis | FlyBodyRotationAxis | WingFlyBodyRotationAxis | str],
        joint_preset: "FlyBodyJointPreset | str | None" = None,
        anatomical_joints: list[FlyBodyAnatomicalJoint] | None = None,
    ) -> None:
        if not (joint_preset is None) ^ (anatomical_joints is None):
            raise ValueError(
                "Skeleton must be initiated from either joint_preset or "
                "anatomical_joints, but not both."
            )

        if joint_preset is not None:
            anatomical_joints = FlyBodyJointPreset(joint_preset).to_joint_list()
        self.anatomical_joints = anatomical_joints

        self.joint_lookup = {(j.parent, j.child): j for j in anatomical_joints}
        self.body_segments = orderedset(
            [seg for nodes in self.joint_lookup.keys() for seg in nodes]
        )
        if isinstance(axis_order, AxisOrder):
            warnings.warn(
                "Using a generic AxisOrder with FlyBodySkeleton; "
                "converting to FlyBodyAxisOrder."
            )
            axis_order = axis_order.to_list_of_str()
        self.axis_order = FlyBodyAxisOrder(axis_order)

    def iter_jointdofs(
        self,
        root: FlyBodyBodySegment | str = "c_thorax",
    ) -> Iterator[FlyBodyJointDOF]:
        """Iterate through joint DOFs in depth-first order starting from the root."""
        if isinstance(root, str):
            root = FlyBodyBodySegment(root)
        tree = self.get_tree()
        for parent, child in tree.dfs_edges(root):
            anatomical_joint = self.joint_lookup[(parent, child)]
            for jointdof in anatomical_joint.iter_dofs(self.axis_order):
                yield jointdof

    def get_actuated_dofs_from_preset(
        self, preset: FlyBodyActuatedDOFPreset | str
    ) -> list[FlyBodyJointDOF]:
        """Given a flybody preset of actuated DoFs, return an explicit list of joints."""
        if isinstance(preset, BaseActuatedDOFPreset):
            preset = FlyBodyActuatedDOFPreset(preset.value)
        else:
            preset = FlyBodyActuatedDOFPreset(preset)
        return preset.filter(list(self.iter_jointdofs()))

get_actuated_dofs_from_preset(preset)

Given a flybody preset of actuated DoFs, return an explicit list of joints.

Source code in src/flygym/flybody/anatomy_flybody.py

def get_actuated_dofs_from_preset(
    self, preset: FlyBodyActuatedDOFPreset | str
) -> list[FlyBodyJointDOF]:
    """Given a flybody preset of actuated DoFs, return an explicit list of joints."""
    if isinstance(preset, BaseActuatedDOFPreset):
        preset = FlyBodyActuatedDOFPreset(preset.value)
    else:
        preset = FlyBodyActuatedDOFPreset(preset)
    return preset.filter(list(self.iter_jointdofs()))

iter_jointdofs(root='c_thorax')

Iterate through joint DOFs in depth-first order starting from the root.

Source code in src/flygym/flybody/anatomy_flybody.py

def iter_jointdofs(
    self,
    root: FlyBodyBodySegment | str = "c_thorax",
) -> Iterator[FlyBodyJointDOF]:
    """Iterate through joint DOFs in depth-first order starting from the root."""
    if isinstance(root, str):
        root = FlyBodyBodySegment(root)
    tree = self.get_tree()
    for parent, child in tree.dfs_edges(root):
        anatomical_joint = self.joint_lookup[(parent, child)]
        for jointdof in anatomical_joint.iter_dofs(self.axis_order):
            yield jointdof

WingFlyBodyAxesSet

Bases: AxesSet

Set of rotation axes for wings using WingFlyBody's axis convention.

Source code in src/flygym/flybody/anatomy_flybody.py

class WingFlyBodyAxesSet(AxesSet):
    """Set of rotation axes for wings using WingFlyBody's axis convention."""

    rotation_axis_class = WingFlyBodyRotationAxis

WingFlyBodyAxisOrder

Bases: BaseAxisOrder, Enum

Axis order enum based on WingFlyBodyRotationAxis.

Source code in src/flygym/flybody/anatomy_flybody.py

class WingFlyBodyAxisOrder(BaseAxisOrder, Enum):
    """Axis order enum based on WingFlyBodyRotationAxis."""

    @classmethod
    def _axis_enum_cls(cls):
        return WingFlyBodyRotationAxis

    PITCH_ROLL_YAW = (
        WingFlyBodyRotationAxis.PITCH,
        WingFlyBodyRotationAxis.ROLL,
        WingFlyBodyRotationAxis.YAW,
    )
    PRY = PITCH_ROLL_YAW
    PITCH_YAW_ROLL = (
        WingFlyBodyRotationAxis.PITCH,
        WingFlyBodyRotationAxis.YAW,
        WingFlyBodyRotationAxis.ROLL,
    )
    PYR = PITCH_YAW_ROLL
    ROLL_PITCH_YAW = (
        WingFlyBodyRotationAxis.ROLL,
        WingFlyBodyRotationAxis.PITCH,
        WingFlyBodyRotationAxis.YAW,
    )
    RPY = ROLL_PITCH_YAW
    ROLL_YAW_PITCH = (
        WingFlyBodyRotationAxis.ROLL,
        WingFlyBodyRotationAxis.YAW,
        WingFlyBodyRotationAxis.PITCH,
    )
    RYP = ROLL_YAW_PITCH
    YAW_PITCH_ROLL = (
        WingFlyBodyRotationAxis.YAW,
        WingFlyBodyRotationAxis.PITCH,
        WingFlyBodyRotationAxis.ROLL,
    )
    YPR = YAW_PITCH_ROLL
    YAW_ROLL_PITCH = (
        WingFlyBodyRotationAxis.YAW,
        WingFlyBodyRotationAxis.ROLL,
        WingFlyBodyRotationAxis.PITCH,
    )
    YRP = YAW_ROLL_PITCH

    DONTCARE = PITCH_ROLL_YAW

WingFlyBodyRotationAxis

Bases: BaseRotationAxis

FlyBody wing axis convention.

Same as FlyBody except pitch/roll are swapped: yaw -> z, pitch -> y, roll -> x.

Source code in src/flygym/flybody/anatomy_flybody.py

class WingFlyBodyRotationAxis(BaseRotationAxis):
    """FlyBody wing axis convention.

    Same as FlyBody except pitch/roll are swapped:
    yaw -> z, pitch -> y, roll -> x.
    """

    PITCH = "pitch"
    P = PITCH
    ROLL = "roll"
    R = ROLL
    YAW = "yaw"
    Y = YAW

    @classmethod
    def _vector_by_axis(cls) -> dict[str, tuple[float, float, float]]:
        return {
            "pitch": (0, 1, 0),
            "roll": (1, 0, 0),
            "yaw": (0, 0, 1),
        }