Source code for Karana.KUtils.MultibodyTUI.graph

from abc import ABC, abstractmethod
from typing import TypeVar, Generic

from Karana.Dynamics import PhysicalBody, Multibody
from Karana.Frame import Frame, FrameContainer

T = TypeVar("T")




class GraphAdapter(ABC, Generic[T]):
    """Common interface for various types of directed graphs

    The purpose of this class is to define a common interface for
    wrapper classes that adapt different graphs, each with their own
    interface, to a single common interface.

    """



    @abstractmethod
    def parents(self, node: T, /) -> list[T]:
        """Get the list of parents for a given node"""




    @abstractmethod
    def children(self, node: T, /) -> list[T]:
        """Get the list of children for a given node"""




    @abstractmethod
    def nodes(self) -> list[T]:
        """Get the list of all nodes in the graph"""




    def descendants(self, node: T, /) -> list[T]:
        """Get the list of all descendants for a given node"""
        # Use a dictionary instead of set to preserve ordering
        visited = {}
        stack = []
        stack.append(node)
        results = []
        # do a dfs traversal
        while stack:
            curr = stack.pop()
            for child in self.children(curr):
                if child not in visited:
                    visited[child] = None
                    stack.append(child)
        return list(visited)




    def ancestors(self, node: T, /) -> list[T]:
        """Get the list of all ancestors for a given node"""
        # Use a dictionary instead of set to preserve ordering
        visited = {}
        stack = []
        stack.append(node)
        results = []
        # do a dfs traversal
        while stack:
            curr = stack.pop()
            for parent in self.parents(curr):
                if parent not in visited:
                    visited[parent] = None
                    stack.append(parent)
        return list(visited)






class BodyGraphAdapter(GraphAdapter[PhysicalBody]):
    """Multibody wrapper implementing the GraphAdapater interface"""

    def __init__(self, multibody: Multibody):
        self.multibody = multibody



    def parents(self, body: PhysicalBody) -> list[PhysicalBody]:
        if body == self.multibody.virtualRoot():
            return []
        parent = body.physicalParentBody()
        return [parent]




    def children(self, body: PhysicalBody) -> list[PhysicalBody]:
        # return body.childBodies()
        return body.multibody().childrenBodies(body)




    def nodes(self) -> list[PhysicalBody]:
        return [self.multibody.virtualRoot()] + self.multibody.physicalBodiesList()






class FrameGraphAdapter(GraphAdapter[Frame]):
    """FrameContainer wrapper implementing the GraphAdapater interface"""

    def __init__(self, frame_container: FrameContainer):
        self.frame_container = frame_container



    def parents(self, frame: Frame) -> list[Frame]:
        parent = frame.parent()
        if parent is None:
            return []
        return [parent]




    def children(self, frame: Frame) -> list[Frame]:
        # Not ideal, but there doesn't seem to be a way to get the
        # children directly
        children = []
        for candidate in self.frame_container.frames():
            if candidate.parent() == frame:
                children.append(candidate)
        return children




    def nodes(self) -> list[Frame]:
        return self.frame_container.frames()