# Copyright (c) 2024-2026 Karana Dynamics Pty Ltd. All rights reserved.
#
# NOTICE TO USER:
#
# This source code and/or documentation (the "Licensed Materials") is
# the confidential and proprietary information of Karana Dynamics Inc.
# Use of these Licensed Materials is governed by the terms and conditions
# of a separate software license agreement between Karana Dynamics and the
# Licensee ("License Agreement"). Unless expressly permitted under that
# agreement, any reproduction, modification, distribution, or disclosure
# of the Licensed Materials, in whole or in part, to any third party
# without the prior written consent of Karana Dynamics is strictly prohibited.
#
# THE LICENSED MATERIALS ARE PROVIDED "AS IS" WITHOUT WARRANTY OF ANY KIND.
# KARANA DYNAMICS DISCLAIMS ALL WARRANTIES, EXPRESS OR IMPLIED, INCLUDING
# BUT NOT LIMITED TO WARRANTIES OF MERCHANTABILITY, NON-INFRINGEMENT, AND
# FITNESS FOR A PARTICULAR PURPOSE.
#
# IN NO EVENT SHALL KARANA DYNAMICS BE LIABLE FOR ANY DAMAGES WHATSOEVER,
# INCLUDING BUT NOT LIMITED TO LOSS OF PROFITS, DATA, OR USE, EVEN IF
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGES, WHETHER IN CONTRACT, TORT,
# OR OTHERWISE ARISING OUT OF OR IN CONNECTION WITH THE LICENSED MATERIALS.
#
# U.S. Government End Users: The Licensed Materials are a "commercial item"
# as defined at 48 C.F.R. 2.101, and are provided to the U.S. Government
# only as a commercial end item under the terms of this license.
#
# Any use of the Licensed Materials in individual or commercial software must
# include, in the user documentation and internal source code comments,
# this Notice, Disclaimer, and U.S. Government Use Provision.
"""Collection of types that represent SOADyn classes as DataStructs.
There are DataStructs that represent the classes (the structure of the Multibody and associated classes) and
their states. For example, MultibodyDS is used to represent the data needed to create a Multibody class and all
the associated bodies, nodes, etc. SubGraphStateDS is used to represent the coordinates of the various subhinges,
force nodes, etc. that make up the SubGraph. Together, they can be used to represent the Multibody at an instance
in time (both the structure and state).
"""
import numpy as np
from typing import (
Generic,
Self,
TypeVar,
cast,
TypeAlias,
ClassVar,
Optional,
Annotated,
Literal,
TYPE_CHECKING,
)
from pathlib import Path
from Karana.KUtils.DataStruct import DataStruct, IdMixin, NestedBaseMixin
from Karana.KUtils.Ktyping import (
Vec3,
Mass,
Length3,
Inertia,
Vec,
normCheck,
Acceleration3,
Mat6n,
npSizeCheck,
)
from Karana.KUtils.Kquantities import convert
from Karana.Math import UnitQuaternion, HomTran, SpatialInertia, RotationMatrix, SpatialVector
from Karana.Math import IntegratorType
from Karana.Math.Integrator_types import IntegratorOptionsDS
from Karana.Dynamics import (
PhysicalBody,
HingeType,
PhysicalSubhinge,
ModelManager,
StatePropagator,
MMSolverType,
SubTree,
SubGraph,
Multibody,
CoordData,
SubhingeType,
LockedSubhinge,
PinSubhinge,
LinearSubhinge,
Linear3Subhinge,
SphericalSubhinge,
ScrewSubhinge,
PhysicalHinge,
HingeBase,
LoopConstraintBase,
LoopConstraintCutJoint,
LoopConstraintConVel,
FramePairHinge,
ConstraintNode,
CoordinateConstraint,
CoordBase,
Node,
ConstraintNode,
BilateralConstraintType,
)
from Karana.Frame import FrameContainer, Frame, PrescribedFrameToFrame
from urchin import Link, Joint, URDF, Inertial, JointLimit
from Karana.Scene.Scene_types import ScenePartSpecDS
from Karana.Scene import (
BoxGeometry,
Color,
CylinderGeometry,
SphereGeometry,
PhysicalMaterial,
AssimpImporter,
PhysicalMaterialInfo,
StaticMeshGeometry,
LAYER_PHYSICAL_GRAPHICS,
LAYER_COLLISION,
)
from Karana.Models import BaseKModel
from Karana.Scene import ProxyScene
from Karana.Core import warn
import os
from contextlib import contextmanager
from abc import abstractmethod
from pydantic import SerializeAsAny
if TYPE_CHECKING:
from Karana.Dynamics import PhysicalModalBody, ModalNodeDeformationProvider
HAVE_MODAL_FLEX = True
try:
# Import only for type hints
from Karana.Dynamics import PhysicalModalBody, ModalNodeDeformationProvider
HAVE_MODAL_FLEX = True
except:
HAVE_MODAL_FLEX = False
SubhingePyType = (
LockedSubhinge
| PinSubhinge
| LinearSubhinge
| Linear3Subhinge
| ScrewSubhinge
| SphericalSubhinge
)
[docs]
class ConstraintFrameDS(DataStruct):
"""DataStruct used to represent a frame that is part of a constraint.
Parameters
----------
name : str
Name of the frame.
translation : Vec3
Translation of the frame relative to its parent.
unit_quaternion : UnitQuaternion
UnitQuaternion that represents the orientation of the frame relative to its parent.
parent : Literal["newtonian","vroot"]
Parent frame. For constraint frames, this can only be the newtonian frame or the vroot (virtual root).
"""
name: str
translation: Length3
unit_quaternion: UnitQuaternion
parent: Literal["newtonian", "vroot"]
@property
def rel_transform(self) -> HomTran:
"""The HomTran of the constraint frame relative to it's parent."""
return HomTran(self.unit_quaternion, convert(self.translation))
[docs]
def toFrame(self, mb: Multibody) -> Frame:
"""Create an instance of a Frame from this ConstraintFrameDS.
Parameters
----------
mb : Multibody
The Multibody to add the constraint frame to.
Returns
-------
Frame
Frame instance created from this ConstraintFrameDS.
"""
fc = mb.getNewtonianFrame().container()
if self.parent == "newtonian":
parent = mb.getNewtonianFrame()
else:
parent = cast(Frame, mb.virtualRoot())
frame = Frame(self.name, fc)
pf2f = PrescribedFrameToFrame(parent, frame)
pf2f.setRelTransform(self.rel_transform)
return frame
[docs]
@classmethod
def fromFrame(cls, frame: Frame, mb: Multibody) -> Self:
"""Create an instance of ConstraintFrameDS from the provided Frame.
Parameters
----------
frame : Frame
The Frame used to create the ConstraintFrameDS.
mb : Multibody
The Multibody using the constraint frame.
Returns
-------
Self
A ConstraintFrameDS that represents the provided Frame.
"""
dark = frame.edge().oframe()
if dark == mb.getNewtonianFrame():
parent = "newtonian"
elif dark == cast(Frame, mb.virtualRoot()):
parent = "vroot"
else:
raise ValueError(
f'Frame "{frame.name()}" is not attached to the multibody\'s newtonian frame or virtual root frame.'
)
T = frame.edge().relTransform()
return cls(
name=frame.name(),
translation=T.getTranslation(),
unit_quaternion=T.getUnitQuaternion(),
parent=parent,
)
[docs]
class NodeDS(DataStruct, IdMixin[Node], NestedBaseMixin):
"""DataStruct used to represent a Node.
Parameters
----------
name : str
The name of the node.
force_node : bool
Whether or not this node is a force node.
translation : Length3 = np.zeros(3)
The body-to-node translation.
unit_quaternion : UnitQuaternion = UnitQuaternion(0.0, 0.0, 0.0, 1.0)
The body-to-node unit quaternion.
"""
name: str
constraint_node: bool
force_node: bool
translation: Length3 = np.zeros(3)
unit_quaternion: UnitQuaternion = UnitQuaternion(0.0, 0.0, 0.0, 1.0)
@property
def body_to_node_transform(self) -> HomTran:
"""A HomTran that represents the body to node transform of the node.
Returns
-------
HomTran
HomTran that represents the body to node transform of the node.
"""
return HomTran(self.unit_quaternion, convert(self.translation))
[docs]
def toNode(self, body: PhysicalBody) -> Node:
"""Create a Node from this NodeDS.
Parameters
----------
body : PhysicalBody
The body to attach the node to.
"""
if self.constraint_node:
nd = ConstraintNode.lookupOrCreate(self.name, body)
else:
nd = Node.lookupOrCreate(self.name, bd=body, force_node=self.force_node)
nd.setBodyToNodeTransform(self.body_to_node_transform)
self.addObjectFromId(nd)
return nd
[docs]
@classmethod
def fromNode(cls, nd: Node) -> Self:
"""Create a NodeDS from the given Node.
Parameters
----------
nd : Node
The Node to use to create this NodeDS.
"""
T = nd.getBodyToNodeTransform()
dark = cls(
name=nd.name(),
constraint_node=isinstance(nd, ConstraintNode),
force_node=nd.isExternalForceNode(),
translation=T.getTranslation(),
unit_quaternion=T.getUnitQuaternion(),
)
dark._id = nd.id()
return dark
[docs]
class ForceNodeStateDS(DataStruct):
"""This class holds state data for a force node.
Parameters
----------
name : str
Name of the node.
external_spatial_force : SpatialVector
The external spatial force on the node in the local node frame.
"""
name: str
external_spatial_force: SpatialVector
[docs]
@classmethod
def fromNode(cls, force_node: Node) -> Self:
"""Create a ForceNodeStateDS from the provided node.
Parameters
----------
force_node : Node
The force node used to create the ForceNodeStateDS.
Returns
-------
Self
The ForceNodeStateDS.
"""
return cls(name=force_node.name(), external_spatial_force=force_node.getSpForce())
[docs]
def toNode(self, nd: Node) -> None:
"""Set the state of the provided node.
Parameters
----------
nd : Node
The node to set the state of.
"""
nd.setExternalSpForce(self.external_spatial_force)
[docs]
class SubhingeDS(DataStruct, NestedBaseMixin):
"""Base class for all subhinges.
Parameters
----------
subhinge_type: SubhingeType
The subhinge type associated with the subhinge.
"""
subhinge_type: SubhingeType
[docs]
@classmethod
def fromSubhinge(cls, sh: SubhingePyType) -> "SubhingeDS":
"""Create an instance of SubhingeDS from the provided Subhinge.
Parameters
----------
sh : SubhingeType
The Subhinge used to create the SubhingeDS.
Returns
-------
Self
A SubhingeDS that represents the provided Subhinge.
"""
if isinstance(sh, PinSubhinge):
return PinSubhingeDS.fromSubhinge(sh)
elif isinstance(sh, LockedSubhinge):
return LockedSubhingeDS.fromSubhinge(sh)
elif isinstance(sh, LinearSubhinge):
return LinearSubhingeDS.fromSubhinge(sh)
elif isinstance(sh, SphericalSubhinge):
return SphericalSubhingeDS.fromSubhinge(sh)
elif isinstance(sh, Linear3Subhinge):
return Linear3SubhingeDS.fromSubhinge(sh)
else: # Screw subhinge
sh = cast(ScrewSubhinge, sh)
return ScrewSubhingeDS.fromSubhinge(sh)
[docs]
@abstractmethod
def setSubhingeValues(self, sh: SubhingePyType) -> None:
"""Set the subhinge values for the provided subhinge."""
...
[docs]
class PinSubhingeDS(SubhingeDS):
"""PinSubhinge DataStruct.
Parameters
----------
unit_axis : Vec3
Unit axis for the pin subhinge.
joint_limits: Annotated[Vec, npSizeCheck((2,))]
The lower and upper bound for the joint limit.
prescribed : bool
Determines if the subhinge is prescribed or not.
"""
subhinge_type: SubhingeType = SubhingeType.PIN
unit_axis: Annotated[Vec3, normCheck(1.0)]
joint_limits: Annotated[Vec, npSizeCheck((2,))] = np.array([-np.inf, np.inf])
prescribed: bool
[docs]
@classmethod
def fromSubhinge(cls, sh: PinSubhinge) -> Self: # pyright: ignore - override is okay.
"""Create a PinSubhingeDS from a PinSubhinge."""
joint_limits = sh.getJointLimits()
if np.all(np.isnan(joint_limits)):
joint_limits = np.array([-np.inf, np.inf])
return cls(
subhinge_type=SubhingeType.PIN,
unit_axis=sh.getUnitAxis(),
prescribed=sh.getPrescribed(),
joint_limits=joint_limits,
)
[docs]
def setSubhingeValues(self, sh: PinSubhinge) -> None: # pyright: ignore - override is okay.
"""Set the values of a PinSubhinge."""
sh.setUnitAxis(self.unit_axis)
sh.setPrescribed(self.prescribed)
if not np.all(self.joint_limits == np.array([-np.inf, np.inf])):
sh.setJointLimits(self.joint_limits)
[docs]
class LinearSubhingeDS(SubhingeDS):
"""LinearSubhinge DataStruct.
Parameters
----------
unit_axis : Vec3
Unit axis for the pin subhinge.
joint_limits: Annotated[Vec, npSizeCheck((2,))]
The lower and upper bound for the joint limit.
prescribed : bool
Determines if the subhinge is prescribed or not.
"""
subhinge_type: SubhingeType = SubhingeType.LINEAR
unit_axis: Annotated[Vec3, normCheck(1.0)]
joint_limits: Annotated[Vec, npSizeCheck((2,))] = np.array([-np.inf, np.inf])
prescribed: bool
[docs]
@classmethod
def fromSubhinge(cls, sh: LinearSubhinge) -> Self: # pyright: ignore - override is okay.
"""Create a LinearSubhingeDS from a LinearSubhinge."""
joint_limits = sh.getJointLimits()
if np.all(np.isnan(joint_limits)):
joint_limits = np.array([-np.inf, np.inf])
return cls(
unit_axis=sh.getUnitAxis(), prescribed=sh.getPrescribed(), joint_limits=joint_limits
)
[docs]
def setSubhingeValues(self, sh: LinearSubhinge) -> None: # pyright: ignore - override is okay.
"""Set the values of a LinearSubhinge."""
sh.setUnitAxis(self.unit_axis)
sh.setPrescribed(self.prescribed)
if not np.all(self.joint_limits == np.array([-np.inf, np.inf])):
sh.setJointLimits(self.joint_limits)
[docs]
class SphericalSubhingeDS(SubhingeDS):
"""SphericalSubhinge DataStruct.
prescribed : bool
Determines if the subhinge is prescribed or not.
"""
subhinge_type: SubhingeType = SubhingeType.SPHERICAL
prescribed: bool
[docs]
@classmethod
def fromSubhinge(cls, sh: SphericalSubhinge) -> Self: # pyright: ignore - override is okay.
"""Create a SphericalSubhingeDS from a SphericalSubhinge."""
return cls(prescribed=sh.getPrescribed())
[docs]
def setSubhingeValues(
self, sh: SphericalSubhinge
) -> None: # pyright: ignore - override is okay.
"""Set the values of a SphericalSubhinge."""
sh.setPrescribed(self.prescribed)
[docs]
class Linear3SubhingeDS(SubhingeDS):
"""Linear3Subhinge DataStruct.
Parameters
----------
prescribed : bool
Determines if the subhinge is prescribed or not.
"""
subhinge_type: SubhingeType = SubhingeType.LINEAR3
prescribed: bool
[docs]
@classmethod
def fromSubhinge(cls, sh: Linear3Subhinge) -> Self: # pyright: ignore - override is okay.
"""Create a Linear3SubhingeDS from a Linear3Subhinge."""
return cls(prescribed=sh.getPrescribed())
[docs]
def setSubhingeValues(self, sh: Linear3Subhinge) -> None: # pyright: ignore - override is okay.
"""Set the values of a Linear3Subhinge."""
sh.setPrescribed(self.prescribed)
[docs]
class LockedSubhingeDS(SubhingeDS):
"""LockedSubhinge DataStruct."""
subhinge_type: SubhingeType = SubhingeType.LOCKED
[docs]
@classmethod
def fromSubhinge(cls, _: LockedSubhinge) -> Self: # pyright: ignore - override is okay.
"""Create a LockedSubhingeDS from a LockedSubhinge."""
return cls()
[docs]
def setSubhingeValues(self, _: LockedSubhinge) -> None: # pyright: ignore - override is okay.
"""Set the values of a LockedSubhinge."""
pass
[docs]
class ScrewSubhingeDS(SubhingeDS):
"""ScrewSubhinge DataStruct.
Parameters
----------
prescribed : bool
Determines if the subhinge is prescribed or not.
unit_axis : Vec3
Unit axis for the pin subhinge.
pitch : float
The pitch of the screw subhinge.
joint_limits: Annotated[Vec, npSizeCheck((2,))]
The lower and upper bound for the joint limit.
"""
subhinge_type: SubhingeType = SubhingeType.SCREW
prescribed: bool
unit_axis: Annotated[Vec3, normCheck(1.0)]
pitch: float
joint_limits: Annotated[Vec, npSizeCheck((2,))] = np.array([-np.inf, np.inf])
[docs]
@classmethod
def fromSubhinge(cls, sh: ScrewSubhinge) -> Self: # pyright: ignore - override is okay.
"""Create a ScrewSubhingeDS from a ScrewSubhinge."""
joint_limits = sh.getJointLimits()
if np.all(np.isnan(joint_limits)):
joint_limits = np.array([-np.inf, np.inf])
return cls(
unit_axis=sh.getUnitAxis(),
pitch=sh.getPitch(),
prescribed=sh.getPrescribed(),
joint_limits=joint_limits,
)
[docs]
def setSubhingeValues(self, sh: ScrewSubhinge) -> None: # pyright: ignore - override is okay.
"""Set the values of a ScrewSubhinge."""
sh.setUnitAxis(self.unit_axis)
sh.setPitch(self.pitch)
sh.setPrescribed(self.prescribed)
if not np.all(self.joint_limits == np.array([-np.inf, np.inf])):
sh.setJointLimits(self.joint_limits)
[docs]
class SubhingeStateDS(DataStruct):
"""This class holds the state of a subhinge.
Parameters
----------
Q : Vec
Positional state of the subhinge.
U : Vec
Velocity state of the subhinge.
Udot : Vec
Acceleration state of the subhinge.
T : Vec
Force state of the subhinge.
"""
Q: Vec
U: Vec
Udot: Vec
T: Vec
[docs]
@classmethod
def fromSubhinge(cls, sh: SubhingePyType) -> Self:
"""Create a SubhingeStateDS from a subhinge.
Parameters
----------
sh : SubhingeBase
The subhinge to create this from.
Returns
-------
SubhingeStateDS
A SubhingeStateDS that reflects the state of the provided subhinge.
"""
return cls(Q=sh.getQ(), U=sh.getU(), Udot=sh.getUdot(), T=sh.getT())
[docs]
def toSubhinge(self, sh: SubhingePyType) -> None:
"""Set the state of the provided subhinge.
Parameters
----------
sh : SubhingeBase
The subhinge whose state will be set using this.
"""
sh.setQ(self.Q)
sh.setU(self.U)
sh.setUdot(self.Udot)
sh.setT(self.T)
[docs]
class HingeDS(DataStruct):
"""Base class for all hinge types.
Hinges that have 0 axes associated with them, e.g., BALL,
use this DataStruct. Classes with more axes (or information) will use derived classes.
Parameters
----------
hinge_type : HingeType
The hinge type associated with the hinge.
subhinges: list[SubhingeDS]
A list of subhinges that make up this hinge.
"""
hinge_type: HingeType
subhinges: list[SerializeAsAny[SubhingeDS]]
[docs]
def toHinge(self, parent_bd: PhysicalBody, child_bd: PhysicalBody) -> PhysicalHinge:
"""Create a Hinge between the provided bodies.
Parameters
----------
parent_bd : PhysicalBody
The parent body for the hinge.
child_bd : PhysicalBody
The child body for the hinge.
Returns
-------
PhysicalHinge
The hinge that was created between the two provided bodies.
"""
if self.hinge_type == HingeType.CUSTOM:
hge = PhysicalHinge(
parent_bd, child_bd, HingeType.CUSTOM, [x.subhinge_type for x in self.subhinges]
)
else:
hge = PhysicalHinge(parent_bd, child_bd, self.hinge_type)
for k in range(hge.nSubhinges()):
self.subhinges[k].setSubhingeValues(
hge.subhinge(k) # pyright: ignore These will be the correct type.
)
return hge
[docs]
@classmethod
def fromHinge(cls, hge: FramePairHinge, htype: HingeType) -> Self:
"""Create a HingeDS from the provided hinge.
Parameters
----------
hge : FramePairHinge
The FramePairHinge used to create the HingeDS.
htype : HingeType
The hinge type associated with the FramePairHinge.
Returns
-------
Self
A HingeDS that represents the provided FramePairHinge.
"""
return cls(
hinge_type=htype,
subhinges=cast(
list[SubhingeDS],
[
SubhingeDS.fromSubhinge(cast(SubhingePyType, hge.subhinge(k)))
for k in range(hge.nSubhinges())
],
),
)
[docs]
class HingeStateDS(DataStruct):
"""Represents the state of a hinge.
Parameters
----------
subhinges : list[SubhingeStateDS]
Data for the hinge's subhinges.
"""
subhinges: list[SubhingeStateDS]
[docs]
@classmethod
def fromHinge(cls, hge: HingeBase) -> Self:
"""Create a HingeStateDS from the provided hinge.
Parameters
----------
hge : HingeBase
The hinge used to this HingeStateDS.
Returns
-------
HingeStateDS
A HingeStateDS that reflects the state of the provided hinge.
"""
return cls(
subhinges=[
SubhingeStateDS.fromSubhinge(cast(SubhingePyType, hge.subhinge(k)))
for k in range(hge.nSubhinges())
]
)
[docs]
def toHinge(self, hge: HingeBase) -> None:
"""Set the state of the provided hinge using this.
Parameters
----------
hge : HingeBase
The hinge whose state will be set from this.
"""
for k, sh in enumerate(self.subhinges):
sh.toSubhinge(cast(SubhingePyType, hge.subhinge(k)))
[docs]
class LoopConstraintDS(DataStruct, NestedBaseMixin):
"""Represents a loop constraint.
Parameters
----------
name: str
The name of the LoopConstraintCutJoint.
get_oframe: dict[Literal["body","frame","node"],str]
A dictionary that specifies how to get to the associated oframe of the LoopConstraint.
This uses body names, frame names, and node names to drill down to the correct location.
get_pframe: dict[Literal["body","frame","node"],str]
A dictionary that specifies how to get to the associated pframe of the LoopConstraint.
This uses body names, frame names, and node names to drill down to the correct location.
"""
name: str
get_oframe: dict[Literal["body", "frame", "node"], str]
get_pframe: dict[Literal["body", "frame", "node"], str]
[docs]
@classmethod
def fromLoopConstraint(cls, lc: LoopConstraintBase) -> Self:
"""Create a LoopConstraintDS from the provided LoopConstraintBase.
Parameters
----------
lc : LoopConstraintBase
The LoopConstraintBase used to build this LoopConstraintDS.
Returns
-------
LoopConstraintDS
A LoopConstraintDS with the provided LoopConstraint information.
"""
if isinstance(lc, LoopConstraintCutJoint):
return cast(Self, LoopConstraintCutJointDS.fromLoopConstraintCutJoint(lc))
if isinstance(lc, LoopConstraintConVel):
return cast(Self, LoopConstraintConVelDS.fromLoopConstraintConVel(lc))
raise NotImplementedError(lc)
[docs]
def toCutJointConstraint(self, mb: Multibody) -> LoopConstraintBase:
"""Create a LoopConstraintBase from the LoopConstraintDS.
Parameters
----------
mb : Multibody
The multibody to add this loop constraint to.
Returns
-------
LoopConstraintBase
The LoopConstraintBase created from this LoopConstraintDS.
"""
if isinstance(self, LoopConstraintCutJointDS):
return self.toCutJointConstraintCutJoint(mb)
if isinstance(self, LoopConstraintConVelDS):
return self.toCutJointConstraintConVel(mb)
raise NotImplementedError(repr(self))
[docs]
class LoopConstraintCutJointDS(LoopConstraintDS):
"""Represents a loop constraint across a hinge.
Parameters
----------
name: str
The name of the LoopConstraintCutJoint.
hinge: HingeDS
The hinge information associated with the LoopConstraintCutJoint.
get_oframe: dict[Literal["body","frame","node"],str]
A dictionary that specifies how to get to the associated oframe of the LoopConstraintCutJoint.
This uses body names, frame names, and node names to drill down to the correct location.
get_pframe: dict[Literal["body","frame","node"],str]
A dictionary that specifies how to get to the associated pframe of the LoopConstraintCutJoint.
This uses body names, frame names, and node names to drill down to the correct location.
"""
loop_constraint_type: Literal["LoopConstraintCutJoint"] = "LoopConstraintCutJoint"
hinge: HingeDS
[docs]
@classmethod
def fromLoopConstraintCutJoint(cls, lc: LoopConstraintCutJoint) -> Self:
"""Create a LoopConstraintCutJointDS from the provided LoopConstraintCutJoint.
Parameters
----------
lc : LoopConstraintCutJoint
The LoopConstraintCutJoint used to build this LoopConstraintDS.
Returns
-------
LoopConstraintDS
A LoopConstraintDS with the provided LoopConstraintCutJoint information.
"""
get_oframe: dict[Literal["body", "frame", "node"], str]
get_pframe: dict[Literal["body", "frame", "node"], str]
if isinstance(of := lc.hinge().oframe(), ConstraintNode):
get_oframe = {"body": of.parentBody().name(), "node": of.name()}
else:
get_oframe = {"frame": of.name()}
if isinstance(pf := lc.errorFrameToFrame().pframe(), ConstraintNode):
get_pframe = {"body": pf.parentBody().name(), "node": pf.name()}
else:
get_pframe = {"frame": pf.name()}
return cls(
name=lc.name(),
get_oframe=get_oframe,
get_pframe=get_pframe,
hinge=HingeDS.fromHinge(lc.hinge(), lc.hinge().hingeType()),
)
[docs]
def toCutJointConstraintCutJoint(self, mb: Multibody) -> LoopConstraintCutJoint:
"""Create a LoopConstraintCutJoint from the LoopConstraintDS.
Parameters
----------
mb : Multibody
The multibody to add this loop constraint to.
Returns
-------
LoopConstraintCutJoint
The LoopConstraintCutJoint created from this LoopConstraintDS.
"""
def _getNodeOrFrame(
d: dict[Literal["body", "frame", "node"], str]
) -> ConstraintNode | Frame:
obj = None
for k, v in d.items():
if k == "body":
obj = mb.getBody(v)
elif k == "node":
obj = obj.getConstraintNode(v) # pyright: ignore
else:
fc = mb.getNewtonianFrame().container()
obj = fc.lookupFrame(v)[0]
if not (isinstance(obj, Frame) or isinstance(obj, ConstraintNode)):
raise ValueError("Had issue retrieving oframe or pframe.")
return obj
of = _getNodeOrFrame(self.get_oframe)
pf = _getNodeOrFrame(self.get_pframe)
if self.hinge.hinge_type == HingeType.CUSTOM:
lc = LoopConstraintCutJoint(
mb,
of.frameToFrame(pf),
self.name,
self.hinge.hinge_type,
[x.subhinge_type for x in self.hinge.subhinges],
)
else:
lc = LoopConstraintCutJoint(
self.name, mb, of.frameToFrame(pf), htype=self.hinge.hinge_type
)
for k in range(lc.hinge().nSubhinges()):
self.hinge.subhinges[k].setSubhingeValues(
lc.hinge().subhinge(k) # pyright: ignore These will be the correct type.
)
return lc
[docs]
class LoopConstraintConVelDS(LoopConstraintDS):
"""Represents a constant velocity loop constraint.
Parameters
----------
name: str
The name of the LoopConstraintConVel.
get_oframe: dict[Literal["body","frame","node"],str]
A dictionary that specifies how to get to the associated oframe of the LoopConstraintConVel.
This uses body names, frame names, and node names to drill down to the correct location.
get_pframe: dict[Literal["body","frame","node"],str]
A dictionary that specifies how to get to the associated pframe of the LoopConstraintConVel.
This uses body names, frame names, and node names to drill down to the correct location.
unit_axis: Annotated[Vec3, normCheck(1.0)]
The constraint unit axis
is_rotational: bool
Whether the constraint is rotational or translational
"""
loop_constraint_type: Literal["LoopConstraintConVel"] = "LoopConstraintConVel"
unit_axis: Annotated[Vec3, normCheck(1.0)]
is_rotational: bool
[docs]
@classmethod
def fromLoopConstraintConVel(cls, lc: LoopConstraintConVel) -> Self:
"""Create a LoopConstraintConVelDS from the provided LoopConstraintConVel.
Parameters
----------
lc : LoopConstraintConVel
The LoopConstraintConVel used to build this LoopConstraintConVelDS.
Returns
-------
LoopConstraintConVelDS
A LoopConstraintConVelDS with the provided LoopConstraintConVel information.
"""
get_oframe = {
"body": lc.sourceNode().parentBody().name(),
"node": lc.sourceNode().name(),
}
get_pframe = {
"body": lc.targetNode().parentBody().name(),
"node": lc.targetNode().name(),
}
unit_axis = lc.getUnitAxis()
is_rotational = lc.isRotational()
return cls(
name=lc.name(),
get_oframe=get_oframe, # pyright: ignore - keys are okay
get_pframe=get_pframe, # pyright: ignore - keys are okay
unit_axis=unit_axis,
is_rotational=is_rotational,
)
[docs]
def toCutJointConstraintConVel(self, mb: Multibody) -> LoopConstraintConVel:
"""Create a LoopConstraintConVel from the LoopConstraintConVelDS.
Parameters
----------
mb : Multibody
The multibody to add this loop constraint to.
Returns
-------
LoopConstraintConVel
The LoopConstraintConVel created from this LoopConstraintConVelDS.
"""
def lookupFrame(
d: dict[Literal["body", "frame", "node"], str], /
) -> ConstraintNode | Frame:
return cast(PhysicalBody, mb.getBody(d["body"])).getConstraintNode(d["node"])
of = lookupFrame(self.get_oframe)
pf = lookupFrame(self.get_pframe)
lc = LoopConstraintConVel(self.name, mb, of.frameToFrame(pf))
lc.setUnitAxis(self.unit_axis, self.is_rotational)
return lc
[docs]
class BodyDS(DataStruct, IdMixin[PhysicalBody], NestedBaseMixin):
"""A DataStruct that represents a PhysicalBody.
Parameters
----------
name : str
The name of the body.
sensor_nodes : list[NodeDS]
A list of sensor nodes attached to the body.
force_nodes : list[NodeDS] = []
A list of force nodes attached to the body.
hinge: HingeDS
The HingeDS that describes the hinge.
mass : Mass
The mass of the body.
body_to_cm : Length3 = np.zeros(3)
The body-to-center-of-mass vector.
inertia : Inertia = np.zeros(3,3)
The inertia of the body, given in the body frame.
body_to_joint : Length3 = np.zeros(3)
The body-to-joint translation.
body_to_joint_quat : UnitQuaternion = UnitQuaternion(0.0, 0.0, 0.0, 1.0)
The body-to-joint unit quaternion.
inb_to_joint : Length3 = np.zeros(3)
The inboard-body-to-joint translation.
inb_to_joint_quat : UnitQuaternion = UnitQuaternion(0.0, 0.0, 0.0, 1.0)
The inboard-body-to-joint unit quaternion.
inb_nodal_matrix : Optional[Mat6n] = None
6xn matrix for the inboard body nodal matrix. Should be set if the inboard body is a
PhysicalModalBody. Should be None otherwise.
scene_parts: List[ScenePartSpecDS]
List of import parts that make up the geometry.
inb_nodal_matrix: Optional[Mat6n] = None
Nodal matrix for the onode of the parent body if the parent body is a PhysicalModalBody.
"""
name: str
sensor_nodes: list[SerializeAsAny[NodeDS]] = []
force_nodes: list[SerializeAsAny[NodeDS]] = []
constraint_nodes: list[SerializeAsAny[NodeDS]] = []
hinge: HingeDS
mass: Mass
body_to_cm: Length3 = np.zeros(3)
inertia: Inertia = np.zeros((3, 3))
body_to_joint: Length3 = np.zeros(3)
body_to_joint_quat: UnitQuaternion = UnitQuaternion(0.0, 0.0, 0.0, 1.0)
inb_to_joint: Length3 = np.zeros(3)
inb_to_joint_quat: UnitQuaternion = UnitQuaternion(0.0, 0.0, 0.0, 1.0)
scene_parts: list[ScenePartSpecDS] = []
inb_nodal_matrix: Optional[Mat6n] = None
@property
def body_to_joint_transform(self) -> HomTran:
"""The HomTran between the body frame and the pnode of the hinge."""
return HomTran(self.body_to_joint_quat, convert(self.body_to_joint))
@property
def inb_to_joint_transform(self) -> HomTran:
"""The HomTran between the parent body frame and the onode of the hinge."""
return HomTran(self.inb_to_joint_quat, convert(self.inb_to_joint))
@property
def spatial_inertia(self) -> SpatialInertia:
"""The spatial inertia of the body."""
mass = cast(float, convert(self.mass))
return SpatialInertia(mass, convert(self.body_to_cm), convert(self.inertia))
[docs]
@classmethod
def fromBody(cls, body: PhysicalBody) -> Self:
"""Create a BodyDS from a PhysicalBody.
Parameters
----------
body : PhysicalBody
The PhysicalBody used to create the BodyDS.
Returns
-------
Self
The BodyDS.
"""
spatial_inertia = body.getSpatialInertia()
inb_to_joint_transform = body.onode().getBodyToNodeTransform()
body_to_joint_transform = body.getBodyToJointTransform()
hinge_type = body.parentHinge().hingeType()
if hinge_type is None:
raise ValueError(
f"Cannot create Body DataStruct from '{body.name()}', since it does not have a hinge."
)
hinge_axes = []
if hinge_type.name == "PIN":
from Karana.Dynamics import PinSubhinge
sh = cast(PinSubhinge, body.parentHinge().subhinge(0))
hinge_axes.append(sh.getUnitAxis())
sensor_nodes, force_nodes, constraint_nodes = cls._fromBodyNodes(body)
scene_parts = [ScenePartSpecDS.fromScenePartSpec(sp) for sp in body.getScenePartSpecs()]
htype = body.parentHinge().hingeType()
if htype is None:
raise ValueError(f"Body '{body.name()}' does not have a hinge yet.")
if HAVE_MODAL_FLEX and isinstance(
body.physicalParentBody(),
PhysicalModalBody, # pyright: ignore Won't be unbond if HAVE_MODAL_FLEX is True
):
dp = cast(
ModalNodeDeformationProvider,
body.onode().deformationProvider(), # pyright: ignore Won't be unbond if HAVE_MODAL_FLEX is True
)
inb_nodal_matrix = dp.getNodalMatrix()
else:
inb_nodal_matrix = None
dark = cls(
name=body.name(),
hinge=HingeDS.fromHinge(cast(PhysicalHinge, body.parentHinge()), htype),
mass=spatial_inertia.mass(),
body_to_cm=spatial_inertia.bodyToCm(),
inertia=spatial_inertia.inertia(),
body_to_joint=body_to_joint_transform.getTranslation(),
body_to_joint_quat=body_to_joint_transform.getUnitQuaternion(),
inb_to_joint=inb_to_joint_transform.getTranslation(),
inb_to_joint_quat=inb_to_joint_transform.getUnitQuaternion(),
force_nodes=force_nodes,
sensor_nodes=sensor_nodes,
constraint_nodes=constraint_nodes,
scene_parts=scene_parts,
inb_nodal_matrix=inb_nodal_matrix,
)
dark._id = body.id()
return dark
@classmethod
def _fromBodyNodes(cls, body: PhysicalBody) -> tuple[list[NodeDS], list[NodeDS], list[NodeDS]]:
"""Retrieve nodes from the provided body.
This helper method is defined so that derived classes like ModalBodyDS
can update to use different node types if applicable.
Parameters
----------
body : PhysicalBody
The body to extract node information from.
Returns
-------
tuple[list[NodeDS], list[NodeDS], list[NodeDS]]
Three lists of NodeDSs. These are the sensor nodes, force nodes, and constraint nodes associated
with the body.
"""
sensor_nodes = []
force_nodes = []
for nd in body.nodeList():
if isinstance(nd, ConstraintNode) or nd.isExternalForceNode():
force_nodes.append(NodeDS.fromNode(nd))
else:
sensor_nodes.append(NodeDS.fromNode(nd))
constraint_nodes = [NodeDS.fromNode(nd) for nd in body.constraintNodeList()]
return sensor_nodes, force_nodes, constraint_nodes
[docs]
def toBody(self, parent_body: PhysicalBody) -> PhysicalBody:
"""Create a body from the BodyDS.
Parameters
----------
parent_body: PhysicalBody
Body to attach the new body to.
Returns
-------
PhysicalBody
The new body.
"""
mbody = parent_body.multibody()
new_body = PhysicalBody(self.name, mbody)
self._toBody(parent_body, new_body)
self.addObjectFromId(new_body)
return new_body
def _toBody(self, parent_body: PhysicalBody, new_body: PhysicalBody) -> None:
"""Assign parameters from this BodyDS to a given body.
This helper method is simlar to toBody, but takes in a body argument.
This allows derived classes, e.g., ModalBodyDS, to create a body and pass it
in to have its parameters set.
Parameters
----------
parent_body: PhysicalBody
Body to attach the new body to.
body: PhysicalBody
The body to set the parameters of.
"""
# Create the body
hinge = self.hinge.toHinge(parent_body, new_body)
if hinge.hingeType() != HingeType.FULL6DOF:
new_body.setBodyToJointTransform(self.body_to_joint_transform)
else:
new_body.setBodyToJointTransform(HomTran())
hinge.onode().setBodyToNodeTransform(self.inb_to_joint_transform)
# Set the inb_nodal_matrix if appropriate
if HAVE_MODAL_FLEX and isinstance(
parent_body, PhysicalModalBody # pyright: ignore It is bound if we are here.
):
if self.inb_nodal_matrix is None:
raise ValueError(
f'Tried to set nodal matrix for the onode on "{parent_body.name()}" that connects it to "{new_body.name()}", but the value is None.'
)
dp = cast(
ModalNodeDeformationProvider,
hinge.onode().deformationProvider(), # pyright: ignore It is bound if we are here.
)
dp.setNodalMatrix(self.inb_nodal_matrix)
# Set the spatial inertia
new_body.setSpatialInertia(self.spatial_inertia)
# Add the nodes
for nd in self.sensor_nodes + self.force_nodes + self.constraint_nodes:
nd.toNode(new_body)
# Add the scene geometry
for p in self.scene_parts:
new_body.addScenePartSpec(p.toScenePartSpec())
_hinge_to_joint_map: ClassVar = {
HingeType.LOCKED: "fixed",
HingeType.REVOLUTE: "continuous",
HingeType.SLIDER: "prismatic",
HingeType.FULL6DOF: "floating",
}
_joint_to_hinge_map: ClassVar = {v: k for k, v in _hinge_to_joint_map.items()} | {
"revolute": HingeType.REVOLUTE
}
[docs]
def toUrdf(self, parent_link: str) -> tuple[Link, Joint]:
"""Convert the BodyDS to a urchin Link and Joint. These can be used to create a URDF file.
Parameters
----------
parent_link : str
The name of the parent link for the Joint.
Returns
-------
tuple[Link, Joint]
The Link and Joint associated with this body.
"""
# In URDF, the Link frame (the URDF body frame) is the same as the joint frame.
# This means they have one less frame than we do. This means we need to premultiply by T^-1,
# where T is the body_to_joint_transform to transform things like the inertia.
# See https://wiki.ros.org/urdf/XML/joint for details
body_to_cm_T = HomTran(self.body_to_cm)
joint_to_cm_T = self.body_to_joint_transform.inverse() * body_to_cm_T
i = Inertial(
mass=self.mass,
inertia=self.spatial_inertia.cmInertia(),
origin=joint_to_cm_T.getMatrix(),
)
from urchin import Box, Cylinder, Sphere, Mesh, Material, Visual, Geometry, Collision
visuals = []
collisions = []
for p in self.scene_parts:
# Get the geometry
if isinstance(p.geometry, BoxGeometry):
geom = Geometry(
box=Box(np.array([p.geometry.width, p.geometry.height, p.geometry.depth]))
)
elif isinstance(p.geometry, CylinderGeometry):
geom = Geometry(cylinder=Cylinder(p.geometry.radius, p.geometry.height))
elif isinstance(p.geometry, SphereGeometry):
geom = Geometry(sphere=Sphere(p.geometry.radius))
elif isinstance(p.geometry, StaticMeshGeometry):
if p.geometry.filename != Path():
geom = Geometry(
mesh=Mesh(
str(p.geometry.filename),
False,
scale=p.scale,
lazy_filename=str(p.geometry.filename),
)
)
else:
raise ValueError(
f"StaticMeshGeometry for {p.name} has no filepath. Not sure how to save to a URDF."
)
else:
raise ValueError(f"Geometry of type {p.geometry} is not supported by URDF.")
# Get the material
mat = Material(
p.name + "_material",
color=np.array(
[
p.material.info.color.r(),
p.material.info.color.g(),
p.material.info.color.b(),
p.material.info.color.alpha(),
]
),
)
T = HomTran(p.unit_quaternion, p.translation)
if (p.layers & LAYER_PHYSICAL_GRAPHICS) != 0:
visuals.append(Visual(geom, name=p.name, material=mat, origin=T.getMatrix()))
if (p.layers & LAYER_COLLISION) != 0:
collisions.append(Collision(name=p.name, origin=T.getMatrix(), geometry=geom))
l = Link(name=self.name, inertial=i, visuals=visuals, collisions=collisions)
htype = self.hinge.hinge_type
joint_type = self._hinge_to_joint_map[htype]
if htype == HingeType.SLIDER:
limits = cast(LinearSubhingeDS, self.hinge.subhinges[0]).joint_limits
if np.any(np.isnan(limits)):
jl = JointLimit(effort=np.inf, velocity=np.inf, lower=None, upper=None)
else:
jl = JointLimit(effort=np.inf, velocity=np.inf, lower=limits[0], upper=limits[1])
elif htype == HingeType.REVOLUTE:
limits = cast(PinSubhingeDS, self.hinge.subhinges[0]).joint_limits
if not np.any(np.isnan(limits)):
# Set limits and also change joint type to revolute (will be continuous otherwise)
jl = JointLimit(effort=np.inf, velocity=np.inf, lower=limits[0], upper=limits[1])
joint_type = "revolute"
else:
jl = None
else:
jl = None
j = Joint(
name=f"{parent_link}_to_{self.name}",
joint_type=joint_type,
parent=parent_link,
child=self.name,
origin=self.inb_to_joint_transform.getMatrix(),
limit=jl,
)
if htype in [HingeType.SLIDER, HingeType.REVOLUTE]:
j.axis = self.hinge.subhinges[
0
].unit_axis # pyright: ignore Ignoring, we know this will have the correct type.
if len(self.sensor_nodes) > 0 or len(self.force_nodes) > 0:
warn(
f"There is more than one node on body {self.name}. These will not be transfered over to the URDF."
)
return (l, j)
[docs]
@classmethod
def fromUrdf(cls, link: Link, joint: Joint) -> Self:
"""Convert from a urchin Link and Joint to a BodyDS.
Parameters
----------
link : Link
The URDF Link.
joint : Joint
The URDF parent joint.
Returns
-------
Self
BodyDS for the link and joint.
"""
# The URDF has one less frame than we do, since they always put their body frame at the joint. Hence, we don't know
# where the body frame actually is. What we do instead, is put the body frame at the center of mass. Hence, we set the
# body_to_joint based on the inertia's origin, and body_to_cm is always 0.
hinge_type = cls._joint_to_hinge_map.get(joint.joint_type, None)
if hinge_type is None:
raise ValueError(
f"Got a URDF with joint type {joint.joint_type}, which does not have a complimentary hinge type."
)
if hinge_type == HingeType.REVOLUTE:
joint_limits = np.array([np.nan, np.nan])
if (dark := joint.limit) is not None:
lower = -np.inf if dark.lower is None else dark.lower
upper = np.inf if dark.upper is None else dark.upper
joint_limits = np.array([lower, upper])
hinge = HingeDS(
hinge_type=hinge_type,
subhinges=cast(
list[SubhingeDS],
[
PinSubhingeDS(
unit_axis=joint.axis, prescribed=False, joint_limits=joint_limits
)
],
),
)
elif hinge_type == HingeType.SLIDER:
joint_limits = np.array([np.nan, np.nan])
if (dark := joint.limit) is not None:
lower = -np.inf if dark.lower is None else dark.lower
upper = np.inf if dark.upper is None else dark.upper
joint_limits = np.array([lower, upper])
hinge = HingeDS(
hinge_type=hinge_type,
subhinges=cast(
list[SubhingeDS],
[
LinearSubhingeDS(
unit_axis=joint.axis, prescribed=False, joint_limits=joint_limits
)
],
),
)
elif hinge_type == HingeType.LOCKED:
hinge = HingeDS(
hinge_type=hinge_type,
subhinges=cast(
list[SubhingeDS],
[LockedSubhingeDS()],
),
)
elif hinge_type == HingeType.FULL6DOF:
hinge = HingeDS(
hinge_type=hinge_type,
subhinges=[
Linear3SubhingeDS(prescribed=False),
SphericalSubhingeDS(prescribed=False),
],
)
elif hinge_type == HingeType.LOCKED:
hinge = HingeDS(hinge_type=hinge_type, subhinges=[])
else:
raise ValueError(f"Hinge type {hinge_type} should not exist on URDFs.")
# In URDFs, the link frame is at the pnode location. Thus, we need to get the inerita
# relative to that.
dark_quat = UnitQuaternion(RotationMatrix(link.inertial.origin[0:3, 0:3]))
cm_to_body_T = HomTran(dark_quat, link.inertial.origin[0:3, 3].flatten())
cm_inertia = SpatialInertia(link.inertial.mass, np.zeros(3), link.inertial.inertia)
# Inertia about the link frame
inertia = cm_inertia.parallelAxis(cm_to_body_T)
inb_to_joint_quat = UnitQuaternion(RotationMatrix(joint.origin[0:3, 0:3]))
from urchin import Visual
scene_parts = []
for p in link.visuals + link.collisions:
quat = UnitQuaternion(RotationMatrix(p.origin[0:3, 0:3]))
trans = p.origin[0:3, 3].flatten()
if p.geometry.box is not None:
b = p.geometry.box
geom = BoxGeometry(*b.size)
elif p.geometry.cylinder is not None:
c = p.geometry.cylinder
geom = CylinderGeometry(c.radius, c.length)
elif p.geometry.sphere is not None:
s = p.geometry.sphere
geom = SphereGeometry(s.radius)
else:
if p.geometry.mesh.filename:
importer = AssimpImporter()
res = importer.importFrom(p.geometry.mesh.filename)
for k, part in enumerate(res.parts):
dark = ScenePartSpecDS.fromScenePartSpec(part)
if p.name is None:
import random
tag = int(random.random() * 10000)
p.name = f"unnamed{tag}"
if len(res.parts) == 1:
dark.name = p.name
else:
dark.name = p.name + f"_mesh{k}"
dark.translation = trans
dark.unit_quaternion = quat
dark.layers = (
LAYER_PHYSICAL_GRAPHICS if isinstance(p, Visual) else LAYER_COLLISION
)
scene_parts.append(dark)
continue
else:
raise ValueError(f"Not sure what to do with geometry {p.geometry}")
info = PhysicalMaterialInfo()
if isinstance(p, Visual):
info.color = Color.fromRGBA(*p.material.color)
mat = PhysicalMaterial(info)
scene_parts.append(
ScenePartSpecDS(
name=p.name if p.name is not None else "unnamed",
material=mat,
geometry=geom,
unit_quaternion=quat,
translation=trans,
scale=np.ones(3),
layers=LAYER_PHYSICAL_GRAPHICS if isinstance(p, Visual) else LAYER_COLLISION,
)
)
return cls(
name=link.name,
hinge=hinge,
mass=inertia.mass(),
body_to_cm=inertia.bodyToCm(),
inertia=inertia.inertia(),
body_to_joint=np.zeros(3),
body_to_joint_quat=UnitQuaternion(0, 0, 0, 1),
inb_to_joint=joint.origin[0:3, 3].flatten(),
inb_to_joint_quat=inb_to_joint_quat,
scene_parts=scene_parts,
)
[docs]
class BodyStateDS(DataStruct, NestedBaseMixin):
"""This class holds state data for a body.
Parameters
----------
name : str
The name of the body.
grav_accel : Acceleration3
The acceleration due to gravity.
hinge : HingeStateDS
The state of the hinge.
force_node_states : list[ForceNodeStateDS] = []
The states of the force nodes on the body.
"""
name: str
grav_accel: Acceleration3
hinge: HingeStateDS
force_node_states: list[ForceNodeStateDS] = []
[docs]
@classmethod
def fromBody(cls, bd: PhysicalBody) -> Self:
"""Create a BodyStateDS from the provided body.
Parameters
----------
bd: PhysicalBody
The body used to create the state.
Returns
-------
Self
The BodyStateDS.
"""
return cls(
name=bd.name(),
grav_accel=bd.getGravAccel(),
hinge=HingeStateDS.fromHinge(bd.parentHinge()),
force_node_states=[
ForceNodeStateDS.fromNode(nd)
for nd in bd.nodeList()
if nd.isExternalForceNode() or isinstance(nd, ConstraintNode)
],
)
[docs]
def toBody(self, bd: PhysicalBody) -> None:
"""Set the state of the provided body.
Parameters
----------
bd : PhysicalBody
The body to set the state of.
"""
bd.setGravAccel(convert(self.grav_accel))
self.hinge.toHinge(bd.parentHinge())
for fn in self.force_node_states:
# skip if the force node does not exist (may be since the
# node may be contact node created on the fly - so not part
# of the original model)
nd = bd.getNode(fn.name)
if nd:
fn.toNode(nd)
else:
print(f"WARNING: Unable to find {fn.name} force node - skipping")
LinkJointTree: TypeAlias = list[tuple[Link, Joint, "LinkJointTree"]]
from Karana.Dynamics._body_ds_type import body_dict, body_state_dict
[docs]
class BodyWithContextDS(DataStruct):
"""A BodyDS with context, i.e., with children that are defined by `BodyDS`s.
Parameters
----------
body : BodyDS
The body we are defining with context.
children : list[BodyWithContextDS]
The child bodies of this body, also given as `BodyWithContextDS`s.
"""
body: SerializeAsAny[BodyDS]
children: list["BodyWithContextDS"] = []
[docs]
def toBody(self, parent_body: PhysicalBody) -> PhysicalBody:
"""Create a body from the BodyWithContextDS.
Parameters
----------
parent_body: PhysicalBody
Body to attach the new body to.
Returns
-------
PhysicalBody
The new body.
"""
body = self.body.toBody(parent_body)
for c in self.children:
c.toBody(body)
return body
[docs]
@classmethod
def fromBody(cls, body: PhysicalBody, st: SubTree) -> Self:
"""Create a BodyWithContextDS from a PhysicalBody.
Parameters
----------
body : PhysicalBody
The PhysicalBody used to create the BodyWithContextDS.
st : SubTree
The SubTree to use for context (the body's children in the bu-tree may be a subset of the full set of physical children bodies
Returns
-------
Self
The BodyWithContextDS.
"""
children = [BodyWithContextDS.fromBody(x, st) for x in st.childrenBodies(body)]
body_ds = None
for k, v in body_dict.items():
if isinstance(body, k):
body_ds = v(body)
break
if body_ds is None:
body_ds = BodyDS.fromBody(body)
return cls(body=cast(BodyDS, body_ds), children=children)
[docs]
def toUrdf(self, parent_link: str) -> tuple[list[Link], list[Joint]]:
"""Convert this BodyWithContextDS to a list of Links and Joints. These can be used to create a URDF file.
Parameters
----------
parent_link : str
The name of the parent link for the body of this BodyWithContextDS.
Returns
-------
tuple[list[Link], list[Joint]]
The Links and Joints associated with the body and children of this BodyWithContextDS.
"""
links = []
joints = []
# Add the body first
l, j = self.body.toUrdf(parent_link)
links.append(l)
joints.append(j)
# Now, add the children bodies
for c in self.children:
c_links, c_joints = c.toUrdf(self.body.name)
links += c_links
joints += c_joints
return links, joints
[docs]
@classmethod
def fromUrdf(cls, link: Link, joint: Joint, children_urdf: LinkJointTree):
"""Convert from links and joints to BodyWithContextDS.
Parameters
----------
link : Link
The link to use for creating the body.
joint : Joint
The joint to use for creating the body.
children_urdf : LinkJointTree
A LinkJointTree that defines the links and joints to use for the
children. This is an arbitrarily nested list of tuples of links
and joints, each of which can be used to define a body.
Returns
-------
Self
An instance of BodyWithContextDS.
"""
body = BodyDS.fromUrdf(link, joint)
children = []
for child in children_urdf:
children.append(BodyWithContextDS.fromUrdf(*child))
return cls(body=body, children=children)
[docs]
class BodyStateWithContextDS(DataStruct):
"""A BodyStateDS with context, i.e., with children that are defined by `BodyStateDS`s.
Parameters
----------
body : BodyStateDS
The body state we are defining with context.
children : list[BodyStateWithContextDS]
The state of child bodies of this body, also given as `BodyStateWithContextDS`s.
Returns
-------
Self
The BodyStateWithContextDS.
"""
body: SerializeAsAny[BodyStateDS]
children: list["BodyStateWithContextDS"] = []
[docs]
@classmethod
def fromBody(cls, bd: PhysicalBody, st: SubTree) -> Self:
"""Create a BodyStateWithContextDS from the provided PhysicalBody.
This will capture the provided PhysicalBody state as a BodyStateDS
in the body field and all of its children as BodyStateWithContextDSs
in the children field. This happens recursively.
Parameters
----------
bd : PhysicalBody
The body used to create the BodyStateWithContextDS.
st : SubTree
The SubTree to use for context (the body's children in the bu-tree may be a subset of the full set of physical children bodies
Returns
-------
Self
An instance of BodyStateWithContextDS that represents the provided PhysicalBody.
"""
body_state_ds = None
for k, v in body_state_dict.items():
if isinstance(bd, k):
body_state_ds = v(bd)
break
if body_state_ds is None:
body_state_ds = BodyStateDS.fromBody(bd)
return cls(
body=cast(BodyStateDS, body_state_ds),
children=[BodyStateWithContextDS.fromBody(x, st) for x in st.childrenBodies(bd)],
)
[docs]
def toBody(self, bd: PhysicalBody) -> None:
"""Set the state of the provided body and its children.
Parameters
----------
bd : PhysicalBody
The body to set the state of.
"""
self.body.toBody(bd)
# don't worry about sub-tree context here, since we are checking
# if the body is available or not anyway
cbodies = {c.name(): c for c in bd.multibody().childrenBodies(bd)}
for c in self.children:
cbody = cbodies.get(c.body.name, None)
if cbody is None:
raise ValueError(
f'Cannot find body with name "{c.body.name}" attached to body with name "{bd.name()}".'
)
c.toBody(cbody)
[docs]
class CoordinateConstraintDS(DataStruct, IdMixin[CoordinateConstraint]):
"""CoordinateConstraint DataStruct.
Parameters
----------
name : str
Name of the CoordinateConstraint.
obody_name: str
Name of the body whose parent hinge the osubhinge belongs to
osubhinge_index: int
Index of the osubhinge within its hinge
pbody_name: str
Name of the body whose parent hinge the psubhinge belongs to
psubhinge_index: int
Index of the psubhinge within its hinge
scale_ratio: float
The coordinate scale ratio
"""
name: str
obody_name: str
osubhinge_index: int
pbody_name: str
psubhinge_index: int
scale_ratio: float
[docs]
@classmethod
def fromCoordinateConstraint(cls, cc: CoordinateConstraint) -> Self:
"""Create a CoordinateConstraintDS from a CoordinateConstraint.
Parameters
----------
cc : CoordinateConstraint
The CoordinateConstraint used to create the CoordinateConstraintDS.
Returns
-------
Self
The CoordinateConstraintDS.
"""
osubhinge = cc.osubhinge()
psubhinge = cc.psubhinge()
return cls(
name=cc.name(),
obody_name=cast(PhysicalHinge, osubhinge.parentHinge()).pnode().parentBody().name(),
osubhinge_index=osubhinge.getIndex(),
pbody_name=cast(PhysicalHinge, psubhinge.parentHinge()).pnode().parentBody().name(),
psubhinge_index=psubhinge.getIndex(),
scale_ratio=cc.getScaleRatio(),
)
[docs]
def toCoordinateConstraint(self, mbody: Multibody) -> CoordinateConstraint:
"""Create a CoordinateConstraint from the CoordinateConstraintDS.
Parameters
----------
mbody : Multibody
Multibody the new CoordinateConstraint will belong to
Returns
-------
CoordinateConstraint
The CoordinateConstraint.
"""
obody = mbody.getBody(self.obody_name)
osubhinge = cast(PhysicalSubhinge, obody.parentHinge().subhinge(self.osubhinge_index))
pbody = mbody.getBody(self.pbody_name)
psubhinge = cast(PhysicalSubhinge, pbody.parentHinge().subhinge(self.psubhinge_index))
cc = CoordinateConstraint(self.name, mbody, osubhinge, psubhinge)
cc.setScaleRatio(self.scale_ratio)
return cc
[docs]
class MultibodyDS(DataStruct, IdMixin[Multibody]):
"""Multibody DataStruct.
Parameters
----------
name : str
Name of the multibody.
base_bodies : list[BodyWithContextDS]
The bodies attached to the virtual root.
loop_constraints: list[LoopConstraintDS]
Loop constraints in the multibody
constraint_frames: list[ConstraintFrameDS]
Extra frames that are part of constraints
coordinate_constraints: list[CoordinateConstraintDS]
Coordinate constraints in the multibody
"""
name: str
base_bodies: list[BodyWithContextDS]
loop_constraints: list[SerializeAsAny[LoopConstraintDS]] = []
constraint_frames: list[ConstraintFrameDS] = []
coordinate_constraints: list[CoordinateConstraintDS] = []
[docs]
def toMultibody(
self,
fc: FrameContainer,
newtonian_frame: Optional[Frame] = None,
scene: Optional[ProxyScene] = None,
) -> Multibody:
"""Create a Multibody from this MultibodyDS.
Parameters
----------
fc: FrameContainer
FrameContainer to use to create the multibody.
Returns
-------
Multibody
The new multibody.
"""
# Create the empty multibody
mbody = Multibody(
self.name, fc, newtonian_frame=newtonian_frame
) # pyright: ignore (not picking up that None is okay)
if scene is not None:
mbody.setScene(scene)
# Attach bodies to the virtualRoot
vroot = cast(PhysicalBody, mbody.virtualRoot())
for body in self.base_bodies:
body.toBody(vroot)
# Create any frames that are needed
for f in self.constraint_frames:
f.toFrame(mbody)
# Create loop constraints
for lc in self.loop_constraints:
lc.toCutJointConstraint(mbody)
for cc in self.coordinate_constraints:
cc.toCoordinateConstraint(mbody)
self.addObjectFromId(mbody)
mbody.ensureHealthy()
return mbody
[docs]
def attachToBody(self, root_bd: PhysicalBody, prefix: str = "") -> list[PhysicalBody]:
"""Create the SubGraph of bodies in this DS and attach it to the specified root body.
Parameters
----------
root_bd: PhysicalBody
The root body to attach to.
Returns
-------
list[PhysicalBody]
List of new base_bodies for the new bodies
"""
def _recurseRename(b: BodyWithContextDS, prefix: str):
"""Rename bodies with prefix."""
b.body.name = prefix + b.body.name
for c in b.children:
_recurseRename(c, prefix)
# Attach bodies to the specified body
new_bodies = []
for b in self.base_bodies:
if prefix:
_recurseRename(b, prefix)
new_bodies.append(b.toBody(root_bd))
# Create any frames that are needed
mbody = root_bd.multibody()
for f in self.constraint_frames:
f.toFrame(mbody)
# Create loop constraints
for lc in self.loop_constraints:
lc.toCutJointConstraint(mbody)
for cc in self.coordinate_constraints:
cc.toCoordinateConstraint(mbody)
return new_bodies
[docs]
@classmethod
def fromSubTree(cls, st: SubTree) -> Self:
"""Create a MultibodyDS from a SubTree. This will include only the bodies of the provided SubTree.
Parameters
----------
st : SubTree
The SubTree used to create the MultibodyDS.
Returns
-------
Self
The MultibodyDS.
"""
# vroot = cast(PhysicalBody, st.virtualRoot())
# base_bodies = [BodyWithContextDS.fromBody(x) for x in vroot.childBodies()]
base_bodies = [BodyWithContextDS.fromBody(x, st) for x in st.baseBodies()]
return cls(
name=st.name(),
base_bodies=base_bodies,
)
[docs]
@classmethod
def fromMultibody(cls, mb: Multibody) -> Self:
"""Create a MultibodyDS from a Multibody.
Parameters
----------
mb : Multibody
The Multibody used to create the MultibodyDS.
Returns
-------
Self
The MultibodyDS.
"""
return cls.fromSubGraph(mb)
[docs]
@classmethod
def fromSubGraph(cls, sg: SubGraph) -> Self:
"""Create a MultibodyDS from a SubGraph.
This will be a MultibodyDS with only the bodies and constraints included in the SubGraph.
Parameters
----------
sg : SubGraph
The SubGraph used to create the MultibodyDS.
Returns
-------
Self
The MultibodyDS.
"""
mb = sg.multibody()
# get the base bodies for the SubGraph. When doing so, pass the
# SubGraph as context since the number of children bodies for a
# body in the SubGraph may be a subset of its full set of bodies
base_bodies = [BodyWithContextDS.fromBody(x, sg) for x in sg.baseBodies()]
constraint_frames = []
loop_constraints = []
for lc in sg.enabledConstraints():
if lc.type() in [
BilateralConstraintType.CUTJOINT_LOOP,
BilateralConstraintType.CONVEL_LOOP,
]:
loop_constraints.append(LoopConstraintDS.fromLoopConstraint(lc))
# if not isinstance(of := lc.hinge().oframe(), ConstraintNode):
cf2f = lc.constraintFrameToFrame()
snd = lc.sourceNode()
if not snd:
constraint_frames.append(ConstraintFrameDS.fromFrame(cf2f.oframe(), mb))
tnd = lc.targetNode()
# if not isinstance(pf := lc.errorFrameToFrame().pframe(), ConstraintNode):
if not tnd:
# constraint_frames.append(ConstraintFrameDS.fromFrame(pf, mb))
constraint_frames.append(ConstraintFrameDS.fromFrame(cf2f.pframe(), mb))
coordinate_constraints = [
CoordinateConstraintDS.fromCoordinateConstraint(cc)
for cc in sg.enabledConstraints()
if cc.type() == BilateralConstraintType.COORDINATE
]
dark = cls(
name=sg.name(),
base_bodies=base_bodies,
loop_constraints=loop_constraints,
constraint_frames=constraint_frames,
coordinate_constraints=coordinate_constraints,
)
dark._id = sg.id()
return dark
[docs]
def toUrdf(self, file: Path | str):
"""Convert the MultibodyDS to a URDF file.
Parameters
----------
file: Path | str
The file to write the URDF to.
"""
if isinstance(file, str):
file = Path(file)
# Create a dummy root body to attach everything to.
links = [
Link(
name="root",
inertial=Inertial(mass=1.0, inertia=np.eye(3)),
visuals=None,
collisions=None,
)
]
joints = []
# Add all the links and joints for the base_bodies and their children
for b in self.base_bodies:
b_links, b_joints = b.toUrdf("root")
links += b_links
joints += b_joints
# Print everything to a URDF
urdf = URDF(name=self.name, links=links, joints=joints)
urdf.save(file)
[docs]
@classmethod
def fromUrdf(cls, file: Path | str) -> Self:
"""Read a URDF file and convert it to a MultibodyDS.
Parameters
----------
file : Path | str
The file to write the URDF from.
Returns
-------
MultibodyDS
A MultibodyDS that corresponds to the URDF.
"""
@contextmanager
def changeDir(dir: Path):
original_dir = os.getcwd()
try:
os.chdir(dir)
yield
finally:
os.chdir(original_dir)
# Convert to Path if it is a string
if isinstance(file, str):
file = Path(file)
dir = file.parent
name = file.name
# URDFs have mesh files relative to the URDF file itself. Therefore, we want to
# change directories to that URDF file so mesh loading works appropriately.
with changeDir(dir):
urdf = URDF.load(name, lazy_load_meshes=True)
# Get the root node in the graph
root = urdf.base_link
if root is None:
raise ValueError("Could not get the root of the URDF.")
root = cast(Link, root)
# Use the graph to create link-joint pairs
# The top-level link-joint pairs are what we want to be base bodies
# in the MultibodyDS
def _recurseLink(parent: Link) -> LinkJointTree:
"""Create a LinkJointTree from a parent node.
This is done recusively.
"""
link_joint_tree: LinkJointTree = []
for child in urdf._G.predecessors(parent):
child = cast(Link, child)
grandchildren = _recurseLink(child)
joint: Joint = urdf._G.get_edge_data(child, parent)["joint"]
link_joint_tree.append((child, joint, grandchildren))
return link_joint_tree
link_joint_tree = _recurseLink(root)
return cls(
name=urdf.name,
base_bodies=[BodyWithContextDS.fromUrdf(*x) for x in link_joint_tree],
)
[docs]
class BaseDataDS(DataStruct):
"""DataStruct for BaseData.
Parameters
----------
Q : Vec
Position-level coordinates.
U : Vec
Velocity-level coordinates.
Udot : Vec
Acceleration-level coordinates.
T : Vec
Force-level coordinates.
"""
Q: Vec
U: Vec
Udot: Vec
T: Vec
[docs]
def toCoordBase(self, b: CoordBase):
"""Set the state of the provided CoordBase using this BaseDataDS.
Parameters
----------
b : CoordBase
The CoordBase whose values will be set.
"""
b.setQ(self.Q)
b.setU(self.U)
b.setUdot(self.Udot)
b.setT(self.T)
[docs]
class CoordDataDS(DataStruct):
"""DataStruct for CoordData.
Parameters
----------
base_data : dict[str, BaseDataDS]
Data from the BaseData.
"""
base_data: dict[str, BaseDataDS]
[docs]
@classmethod
def fromCoordData(cls, c: CoordData) -> Self:
"""Create a CoordDataDS that represents the provided CoordData.
Parameters
----------
c : CoordData
The CoordData to represent as a CoordDataDS.
Returns
-------
An instance of CoordDataDS that represents the provided CoordData.
"""
return cls(
base_data={
k: BaseDataDS(Q=v[0], U=v[1], Udot=v[2], T=v[3]) for k, v in c.toMap().items()
}
)
[docs]
def toCoordData(self, c: CoordData) -> None:
"""Set the state of the provided `CoordData`.
Parameters
----------
c : CoordData
The CoordData to set the state of.
"""
for b in c.coordBases():
b = cast(CoordBase, b)
self.base_data[
b.name() # pyright: ignore - All the CoordBases we use should have a name
].toCoordBase(b)
[docs]
class SubTreeStateDS(DataStruct):
"""SubTree State DataStruct.
Parameters
----------
base_bodies : list[BodyStateWithContextDS]
The state of the bodies attached to the virtual root.
"""
base_bodies: list[BodyStateWithContextDS]
[docs]
@classmethod
def fromSubTree(cls, st: SubTree) -> Self:
"""Create a SubTreeStateDS from a SubTree.
Parameters
----------
st : SubTree
The SubTree used to create the SubTreeStateDS.
Returns
-------
Self
The SubTreeStateDS.
"""
# vroot = cast(PhysicalBody, st.virtualRoot())
all_bodies = st.sortedPhysicalBodiesList()
return cls(
base_bodies=[
# BodyStateWithContextDS.fromBody(x) for x in vroot.childBodies() if x in all_bodies
BodyStateWithContextDS.fromBody(x, st)
for x in st.baseBodies()
if x in all_bodies
]
)
[docs]
def toSubTree(self, st: SubTree) -> None:
"""Set the state of the provided multibody.
Parameters
----------
st : SubTree
The multibody to set the state of.
"""
# vroot = cast(PhysicalBody, st.virtualRoot())
# cbodies = {c.name(): c for c in vroot.childBodies()}
cbodies = {c.name(): c for c in st.baseBodies()}
for c in self.base_bodies:
cbody = cbodies.get(c.body.name, None)
if cbody is None:
raise ValueError(
f'Cannot find body with name "{c.body.name}" attached to the virtual root.'
)
c.toBody(cbody)
[docs]
class SubGraphStateDS(DataStruct):
"""SubGraph State DataStruct.
Parameters
----------
sub_tree_state : SubTreeStateDS
The state of the subtree.
constraint_coord_data: CoordDataDS
The state of the constraint coordinate data.
"""
sub_tree_state: SubTreeStateDS
constraint_coord_data: CoordDataDS
[docs]
@classmethod
def fromSubGraph(cls, sg: SubGraph) -> Self:
"""Create a SubGraphStateDS from a SubGraph.
Parameters
----------
sg : SubGraph
The SubGraph used to create the SubGraphStateDS.
Returns
-------
Self
The SubGraphStateDS.
"""
return cls(
sub_tree_state=SubTreeStateDS.fromSubTree(sg),
constraint_coord_data=CoordDataDS.fromCoordData(sg.cutjointCoordData()),
)
[docs]
def toSubGraph(self, sg: SubGraph) -> None:
"""Set the state of the provided multibody.
Parameters
----------
sg : SubGraph
The multibody to set the state of.
"""
self.sub_tree_state.toSubTree(sg)
self.constraint_coord_data.toCoordData(sg.cutjointCoordData())
T = TypeVar("T", bound=BaseKModel)
[docs]
class KModelDS(DataStruct, NestedBaseMixin, IdMixin[T], Generic[T]):
"""A base class for representing a KModel as a DataStruct."""
name: str
period: float | np.timedelta64
[docs]
@abstractmethod
def toModel(self, mm: ModelManager) -> BaseKModel:
"""Create an model instance from this KModelDS.
Parameters
----------
mm : ModelManager
The ModelManager to associate the newly created model instance with.
Returns
-------
BaseKModel
A model instance created from this KModelDS.
"""
raise NotImplementedError("Derived classes must implement toModel")
[docs]
class StatePropagatorDS(DataStruct):
"""A DataStruct that represents the StatePropagator.
Parameters
----------
integrator_type: IntegratorType
Type of integrator used.
integrator_options: Optional[IntegratorOptionsDS] = None
The options associated with the integrator.
solver_type: Optional[MMSolverType] = MMSolverType.UNDEFINED
The options associated with the integrator.
models: list[KModelDS]
The models registered with the StatePropagator.
"""
integrator_type: IntegratorType
integrator_options: Optional[IntegratorOptionsDS] = None
solver_type: Optional[MMSolverType] = MMSolverType.UNDEFINED
models: list[SerializeAsAny[KModelDS]]
[docs]
@classmethod
def fromStatePropagator(cls, sp: StatePropagator) -> Self:
"""Create a StatePropagatorDS from a StatePropagator.
Parameters
----------
body : StatePropagator
The StatePropagator used to create the StatePropagatorDS.
Returns
-------
Self
The StatePropagatorDS.
"""
integ = sp.getIntegrator()
solver_type = sp.solverType()
opts = IntegratorOptionsDS.fromIntegrator(integ)
models: list[KModelDS] = []
for model in sp.getRegisteredModels():
if not IdMixin.objectHasDS(model.id()):
# Only try to create a DS for the model if it was not done already.
# If it was, then this model is already handled by something else, e.g., a Prefab.
if hasattr(model, "toDS"):
models.append(
model.toDS() # pyright: ignore - Will have toDS if we are in this branch
)
else:
warn(
f'Model {model.name()} ({model.typeString()}) does not have a "toDS" method, so it cannot be added to the StatePropagatorDS.'
)
return cls(
integrator_type=integ.getIntegratorType(),
integrator_options=opts,
solver_type=solver_type,
models=models,
)
[docs]
def toStatePropagatorNoModels(self, st: SubTree) -> StatePropagator:
"""Create a StatePropagator from the StatePropagatorDS.
This does NOT add the models to the StatePropagator. If you want the models to, use the
toStatePropagator method.
Parameters
----------
st : SubTree
The SubTree to use to create the StatePropagator.
Returns
-------
StatePropagator
The StatePropagator associated with this StatePropagatorDS, but without any models.
"""
if self.integrator_options is not None:
return StatePropagator(
st,
self.integrator_type,
self.integrator_options.toIntegratorOptions(),
None, # TODO - update to save and restore solver options
self.solver_type,
)
else:
return StatePropagator(st, self.integrator_type)
[docs]
def toStatePropagator(self, st: SubTree) -> StatePropagator:
"""Create a StatePropagator from the StatePropagatorDS.
Parameters
----------
st: SubTree
The SubTree to associate with the StatePropagator.
"""
sp = self.toStatePropagatorNoModels(st)
sp.setTime(0.0)
sp.setState(sp.assembleState())
for model_ds in self.models:
model_ds.toModel(sp)
return sp
