Source code for Karana.KUtils.MultibodyTUI.loopdb

"""Contains the LoopConstraintDatabase class, which is used to query body/loop constriant relationships."""

from Karana.Dynamics import Multibody, PhysicalBody, LoopConstraintBase
from collections import defaultdict
from collections.abc import Mapping, Generator




class LoopConstraintDatabase:
    """Helper to query body / loop constraint relationships."""

    def __init__(self, multibody: Multibody):
        """Create an instance of LoopConstraintDatabase.

        Parameters
        ----------
        multibody: Multibody
            The multibody to do queries within
        """
        self.multibody = multibody

        # map with the list of all the bodies in a loop constraint's loop
        self._lc_related_bodies: Mapping[int, list[PhysicalBody]] = defaultdict(list)

        # map with the list of all the loop constraints that
        # the body in a part of
        self._body_related_lcs: Mapping[int, list[LoopConstraintBase]] = defaultdict(list)

        for lc in self.allLoopConstraints():
            related_bodies = self._computeLoopBodyPath(lc)

            self._lc_related_bodies[lc.id()] = related_bodies
            for body in related_bodies:
                self._body_related_lcs[body.id()].append(lc)



    def allLoopConstraints(self) -> Generator[LoopConstraintBase]:
        """Retrieve all the loop constraints.

        Returns
        -------
        Generator[LoopConstraintBase]
            All the loop constraints as a generator.
        """
        for constraint in self.multibody.enabledConstraints():
            if isinstance(constraint, LoopConstraintBase):
                yield constraint


    def _computeLoopBodyPath(self, lc: LoopConstraintBase) -> list[PhysicalBody]:
        """Compute the body path for the loop constraint.

        This is the alternate path through the body tree connecting the
        source and target nodes of the constraint.
        """
        src_body = lc.sourceNode().parentBody()
        tgt_body = lc.targetNode().parentBody()
        common_ancestor = self.multibody.ancestorBody(src_body, tgt_body)
        src_leg = []
        tgt_leg = []
        while src_body != common_ancestor:
            src_leg.append(src_body)
            src_body = src_body.physicalParentBody()
        while tgt_body != common_ancestor:
            tgt_leg.append(tgt_body)
            tgt_body = tgt_body.physicalParentBody()
        path = src_leg + [common_ancestor] + list(reversed(tgt_leg))
        return path



    def emanatingLoopConstraints(self, body: PhysicalBody) -> list[LoopConstraintBase]:
        """Get all loop constraints directly attached to the body.

        Parameters
        ----------
        body: PhysicalBody
            The given body

        Returns
        -------
        list[LoopConstraintBase]
            The list of loop constraints
        """
        return self.multibody.getBodyLoopConstraints(body)




    def relatedLoopConstraints(self, body: PhysicalBody) -> list[LoopConstraintBase]:
        """Get all loop constraints having the body in their path.

        A loop constraint is included if the path through the body tree
        from their source to target node includes the given body.

        Parameters
        ----------
        body: PhysicalBody
            The given body

        Returns
        -------
        list[LoopConstraintBase]
            The list of loop constraints
        """
        return self._body_related_lcs[body.id()][:]




    def relatedBodies(self, loop_constraint: LoopConstraintBase) -> list[PhysicalBody]:
        """Get all bodies allong the loop constraint's path.

        Lists the bodies along the path from the source to target node
        through the body tree.

        Parameters
        ----------
        loop_constraint: LoopConstraintBase
            The given loop constraint

        Returns
        -------
        list[PhysicalBody]
            The list of bodies
        """
        return self._lc_related_bodies[loop_constraint.id()][:]