openscvx 0.3.2.dev170__py3-none-any.whl

openscvx/symbolic/expr/lie/__init__.py

@@ -0,0 +1,87 @@
+"""Lie algebra operations for rigid body dynamics.
+
+This module provides symbolic expression nodes for Lie algebra operations
+commonly used in 6-DOF rigid body dynamics, robotics, and geometric mechanics.
+These operations enable elegant formulations of Newton-Euler dynamics using
+spatial vectors (twists and wrenches).
+
+The module provides two tiers of functionality:
+
+**Built-in operators** work out of the box and include adjoint/coadjoint
+operators for dynamics (_e.g._ ``Adjoint``, ``AdjointDual``) and frame transformations
+(_e.g._ ``SE3Adjoint``, ``SE3AdjointDual``).
+
+**jaxlie-backed operators** require ``pip install openscvx[lie]`` and provide
+exponential/logarithm maps for SO(3) and SE(3) groups (_e.g._ ``SO3Exp``, ``SO3Log``,
+``SE3Exp``, ``SE3Log``).
+
+Conventions:
+    - Twist (spatial velocity): ξ = [v; ω] where v ∈ ℝ³ is linear velocity
+      and ω ∈ ℝ³ is angular velocity (both in body frame)
+    - Wrench (spatial force): F = [f; τ] where f ∈ ℝ³ is force and τ ∈ ℝ³
+      is torque (both in body frame)
+
+Note:
+    The twist convention [v; ω] (linear first, angular second) matches jaxlie's
+    SE3 tangent parameterization, so no reordering is needed during lowering.
+
+Example:
+    Newton-Euler dynamics for a rigid body using the coadjoint operator::
+
+        import openscvx as ox
+
+        twist = ox.State("twist", shape=(6,))
+        M = ox.Parameter("M", shape=(6, 6), value=spatial_inertia)
+        wrench = ox.Control("wrench", shape=(6,))
+
+        momentum = M @ twist
+        bias_force = ox.lie.AdjointDual(twist, momentum)
+        twist_dot = M_inv @ (wrench - bias_force)
+
+    Product of Exponentials forward kinematics (requires jaxlie)::
+
+        screw_axis = ox.Constant(np.array([0, 0, 0, 0, 0, 1]))
+        theta = ox.State("theta", shape=(1,))
+        T_joint = ox.lie.SE3Exp(screw_axis * theta)  # 4×4 matrix
+
+References:
+    - Murray, Li, Sastry: "A Mathematical Introduction to Robotic Manipulation"
+    - Featherstone: "Rigid Body Dynamics Algorithms"
+    - Sola et al.: "A micro Lie theory for state estimation in robotics"
+"""
+
+# Core operators - no dependencies
+from .adjoint import Adjoint, AdjointDual, SE3Adjoint, SE3AdjointDual
+
+# jaxlie-backed operators - optional dependency
+try:
+    from .se3 import SE3Exp, SE3Log
+    from .so3 import SO3Exp, SO3Log
+
+    _JAXLIE_AVAILABLE = True
+except ImportError:
+    _JAXLIE_AVAILABLE = False
+
+    def _make_stub(name: str):
+        """Create a stub class that raises ImportError on instantiation."""
+
+        def __init__(self, *args, **kwargs):
+            raise ImportError(f"{name} requires jaxlie. Install with: pip install openscvx[lie]")
+
+        return type(name, (), {"__init__": __init__})
+
+    SO3Exp = _make_stub("SO3Exp")
+    SO3Log = _make_stub("SO3Log")
+    SE3Exp = _make_stub("SE3Exp")
+    SE3Log = _make_stub("SE3Log")
+
+__all__ = [
+    "AdjointDual",
+    "Adjoint",
+    "SE3Adjoint",
+    "SE3AdjointDual",
+    "SO3Exp",
+    "SO3Log",
+    "SE3Exp",
+    "SE3Log",
+]

openscvx/symbolic/expr/lie/adjoint.py

@@ -0,0 +1,357 @@
+"""Adjoint and coadjoint operators for rigid body dynamics.
+
+This module provides the core Lie algebra operators for 6-DOF rigid body
+dynamics. These operators require no external dependencies and work with
+the standard openscvx installation.
+
+The module uses the following conventions:
+    - Twist (spatial velocity): ξ = [v; ω] where v ∈ ℝ³ is linear velocity
+      and ω ∈ ℝ³ is angular velocity (both in body frame)
+    - Wrench (spatial force): F = [f; τ] where f ∈ ℝ³ is force and τ ∈ ℝ³
+      is torque (both in body frame)
+    - Momentum: μ = [p; L] where p ∈ ℝ³ is linear momentum and L ∈ ℝ³
+      is angular momentum
+"""
+
+from typing import Tuple
+
+from ..expr import Expr, to_expr
+
+
+class AdjointDual(Expr):
+    """Coadjoint operator ad* for computing Coriolis and centrifugal forces.
+
+    Computes the coadjoint action ad*_ξ(μ) which represents the rate of change
+    of momentum due to body rotation. This is the key term in Newton-Euler
+    dynamics that captures Coriolis and centrifugal effects.
+
+    For se(3), given twist ξ = [v; ω] and momentum μ = [f; τ]:
+
+        ad*_ξ(μ) = [ ω × f + v × τ ]
+                   [     ω × τ     ]
+
+    This appears in the Newton-Euler equations as:
+
+        M @ ξ_dot = F_ext - ad*_ξ(M @ ξ)
+
+    where M is the spatial inertia matrix and F_ext is the external wrench.
+
+    Attributes:
+        twist: 6D twist vector [v; ω] (linear velocity, angular velocity)
+        momentum: 6D momentum vector [p; L] or [f; τ] (linear, angular)
+
+    Example:
+        Compute the bias force (Coriolis + centrifugal) for rigid body dynamics::
+
+            import openscvx as ox
+
+            twist = ox.State("twist", shape=(6,))
+            M = ox.Parameter("M", shape=(6, 6), value=inertia_matrix)
+
+            momentum = M @ twist
+            bias_force = ox.lie.AdjointDual(twist, momentum)
+
+            # In dynamics: twist_dot = M_inv @ (wrench - bias_force)
+
+    Note:
+        The coadjoint is related to the adjoint by: ad*_ξ = -(ad_ξ)^T
+
+        For the special case of pure rotation (v=0) with diagonal inertia,
+        the angular part reduces to the familiar ω × (J @ ω) term.
+
+    See Also:
+        Adjoint: The adjoint operator for twist-on-twist action
+        SSM: 3x3 skew-symmetric matrix for cross products
+    """
+
+    def __init__(self, twist, momentum):
+        """Initialize a coadjoint operator.
+
+        Args:
+            twist: 6D twist vector [v; ω] with shape (6,)
+            momentum: 6D momentum vector [p; L] with shape (6,)
+        """
+        self.twist = to_expr(twist)
+        self.momentum = to_expr(momentum)
+
+    def children(self):
+        return [self.twist, self.momentum]
+
+    def canonicalize(self) -> "Expr":
+        twist = self.twist.canonicalize()
+        momentum = self.momentum.canonicalize()
+        return AdjointDual(twist, momentum)
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Check that inputs are 6D vectors and return output shape.
+
+        Returns:
+            tuple: Shape (6,) for the resulting coadjoint vector
+
+        Raises:
+            ValueError: If twist or momentum do not have shape (6,)
+        """
+        twist_shape = self.twist.check_shape()
+        momentum_shape = self.momentum.check_shape()
+
+        if twist_shape != (6,):
+            raise ValueError(f"AdjointDual expects twist with shape (6,), got {twist_shape}")
+        if momentum_shape != (6,):
+            raise ValueError(f"AdjointDual expects momentum with shape (6,), got {momentum_shape}")
+
+        return (6,)
+
+    def __repr__(self):
+        return f"ad_dual({self.twist!r}, {self.momentum!r})"
+
+
+class Adjoint(Expr):
+    """Adjoint operator ad (Lie bracket) for twist-on-twist action.
+
+    Computes the adjoint action ad_ξ₁(ξ₂) which represents the Lie bracket
+    [ξ₁, ξ₂] of two twists. This is used for velocity propagation in
+    kinematic chains and acceleration computations.
+
+    For se(3), given twists ξ₁ = [v₁; ω₁] and ξ₂ = [v₂; ω₂]:
+
+        ad_ξ₁(ξ₂) = [ξ₁, ξ₂] = [ ω₁ × v₂ - ω₂ × v₁ ]
+                                [     ω₁ × ω₂       ]
+
+    Equivalently using the adjoint matrix:
+
+        ad_ξ = [ [ω]×    0   ]
+               [ [v]×   [ω]× ]
+
+    where [·]× denotes the 3x3 skew-symmetric (cross product) matrix.
+
+    Attributes:
+        twist1: First 6D twist vector [v₁; ω₁]
+        twist2: Second 6D twist vector [v₂; ω₂]
+
+    Example:
+        Compute the Lie bracket of two twists::
+
+            import openscvx as ox
+
+            twist1 = ox.State("twist1", shape=(6,))
+            twist2 = ox.State("twist2", shape=(6,))
+
+            bracket = ox.lie.Adjoint(twist1, twist2)
+
+        Velocity propagation in a kinematic chain::
+
+            # Child link velocity includes parent velocity plus relative motion
+            # V_child = Ad_T @ V_parent + joint_twist * q_dot
+
+    Note:
+        The adjoint satisfies the Jacobi identity and is antisymmetric:
+        ad_ξ₁(ξ₂) = -ad_ξ₂(ξ₁)
+
+    See Also:
+        AdjointDual: The coadjoint operator for momentum dynamics
+    """
+
+    def __init__(self, twist1, twist2):
+        """Initialize an adjoint operator.
+
+        Args:
+            twist1: First 6D twist vector [v; ω] with shape (6,)
+            twist2: Second 6D twist vector [v; ω] with shape (6,)
+        """
+        self.twist1 = to_expr(twist1)
+        self.twist2 = to_expr(twist2)
+
+    def children(self):
+        return [self.twist1, self.twist2]
+
+    def canonicalize(self) -> "Expr":
+        twist1 = self.twist1.canonicalize()
+        twist2 = self.twist2.canonicalize()
+        return Adjoint(twist1, twist2)
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Check that inputs are 6D vectors and return output shape.
+
+        Returns:
+            tuple: Shape (6,) for the resulting Lie bracket
+
+        Raises:
+            ValueError: If either twist does not have shape (6,)
+        """
+        twist1_shape = self.twist1.check_shape()
+        twist2_shape = self.twist2.check_shape()
+
+        if twist1_shape != (6,):
+            raise ValueError(f"Adjoint expects twist1 with shape (6,), got {twist1_shape}")
+        if twist2_shape != (6,):
+            raise ValueError(f"Adjoint expects twist2 with shape (6,), got {twist2_shape}")
+
+        return (6,)
+
+    def __repr__(self):
+        return f"ad({self.twist1!r}, {self.twist2!r})"
+
+
+class SE3Adjoint(Expr):
+    """SE(3) Adjoint representation Ad_T for transforming twists between frames.
+
+    Computes the 6×6 adjoint matrix Ad_T that transforms twists from one
+    coordinate frame to another. Given a transformation T_ab from frame A to
+    frame B, the adjoint transforms a twist expressed in frame A to frame B:
+
+        ξ_b = Ad_{T_ab} @ ξ_a
+
+    For SE(3), given T with rotation R and translation p:
+
+        Ad_T = [ R      0   ]
+               [ [p]×R  R   ]
+
+    where [p]× is the 3×3 skew-symmetric matrix of p.
+
+    This is essential for:
+
+    - Velocity propagation through kinematic chains
+    - Computing geometric Jacobians for manipulators
+    - Recursive Newton-Euler dynamics algorithms
+
+    Attributes:
+        transform: 4×4 homogeneous transformation matrix
+
+    Example:
+        Transform a body twist to the world frame::
+
+            import openscvx as ox
+
+            T_world_body = forward_kinematics(q)  # 4×4 transform
+            twist_body = ox.State("twist_body", shape=(6,))
+
+            # Transform twist to world frame
+            Ad_T = ox.lie.SE3Adjoint(T_world_body)  # 6×6 matrix
+            twist_world = Ad_T @ twist_body
+
+        Compute geometric Jacobian columns::
+
+            # Each column of the geometric Jacobian is Ad_{T_0i} @ ξ_i
+            J_col_i = ox.lie.SE3Adjoint(T_0_to_i) @ screw_axis_i
+
+    Note:
+        The adjoint satisfies: Ad_{T1 @ T2} = Ad_{T1} @ Ad_{T2}
+
+    See Also:
+        - SE3AdjointDual: For transforming wrenches between frames
+        - Adjoint: The small adjoint (Lie bracket) for twist-on-twist action
+    """
+
+    def __init__(self, transform):
+        """Initialize SE3 Adjoint operator.
+
+        Args:
+            transform: 4×4 homogeneous transformation matrix with shape (4, 4)
+        """
+        self.transform = to_expr(transform)
+
+    def children(self):
+        return [self.transform]
+
+    def canonicalize(self) -> "Expr":
+        transform = self.transform.canonicalize()
+        return SE3Adjoint(transform)
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Check that input is a 4×4 matrix and return output shape.
+
+        Returns:
+            tuple: Shape (6, 6) for the adjoint matrix
+
+        Raises:
+            ValueError: If transform does not have shape (4, 4)
+        """
+        transform_shape = self.transform.check_shape()
+        if transform_shape != (4, 4):
+            raise ValueError(
+                f"SE3Adjoint expects transform with shape (4, 4), got {transform_shape}"
+            )
+        return (6, 6)
+
+    def __repr__(self):
+        return f"Ad({self.transform!r})"
+
+
+class SE3AdjointDual(Expr):
+    """SE(3) coadjoint representation Ad*_T for transforming wrenches between frames.
+
+    Computes the 6×6 coadjoint matrix Ad*_T that transforms wrenches from one
+    coordinate frame to another. Given a transformation T_ab from frame A to
+    frame B, the coadjoint transforms a wrench expressed in frame B to frame A:
+
+        F_a = Ad*_{T_ab} @ F_b
+
+    For SE(3), given T with rotation R and translation p:
+
+        Ad*_T = [ R     [p]×R ]
+                [ 0       R   ]
+
+    This is the transpose-inverse of Ad_T: Ad*_T = (Ad_T)^{-T}
+
+    This is essential for:
+
+    - Force/torque propagation in dynamics
+    - Transforming wrenches between end-effector and base frames
+    - Recursive Newton-Euler dynamics algorithms
+
+    Attributes:
+        transform: 4×4 homogeneous transformation matrix
+
+    Example:
+        Transform a wrench from end-effector to base frame::
+
+            import openscvx as ox
+
+            T_base_ee = forward_kinematics(q)  # 4×4 transform
+            wrench_ee = ox.Control("wrench_ee", shape=(6,))
+
+            # Transform wrench to base frame
+            Ad_star_T = ox.lie.SE3AdjointDual(T_base_ee)  # 6×6 matrix
+            wrench_base = Ad_star_T @ wrench_ee
+
+    Note:
+        The coadjoint is related to the adjoint by: Ad*_T = (Ad_T)^{-T}
+
+    See Also:
+        - SE3Adjoint: For transforming twists between frames
+        - AdjointDual: The small coadjoint for Coriolis/centrifugal forces
+    """
+
+    def __init__(self, transform):
+        """Initialize SE3 coadjoint operator.
+
+        Args:
+            transform: 4×4 homogeneous transformation matrix with shape (4, 4)
+        """
+        self.transform = to_expr(transform)
+
+    def children(self):
+        return [self.transform]
+
+    def canonicalize(self) -> "Expr":
+        transform = self.transform.canonicalize()
+        return SE3AdjointDual(transform)
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Check that input is a 4×4 matrix and return output shape.
+
+        Returns:
+            tuple: Shape (6, 6) for the coadjoint matrix
+
+        Raises:
+            ValueError: If transform does not have shape (4, 4)
+        """
+        transform_shape = self.transform.check_shape()
+        if transform_shape != (4, 4):
+            raise ValueError(
+                f"SE3AdjointDual expects transform with shape (4, 4), got {transform_shape}"
+            )
+        return (6, 6)
+
+    def __repr__(self):
+        return f"Ad_dual({self.transform!r})"

openscvx/symbolic/expr/lie/se3.py

@@ -0,0 +1,172 @@
+"""SE(3) Lie group operations for rigid body transformations.
+
+This module provides exponential and logarithm maps for the SE(3) rigid
+transformation group, enabling twist to transformation matrix conversions
+and vice versa. These are essential for Product of Exponentials (PoE)
+forward kinematics in robotic manipulators.
+
+Requires jaxlie: pip install openscvx[lie]
+
+Note:
+    The twist convention [v; ω] (linear first, angular second) matches jaxlie's
+    SE3 tangent parameterization, so no reordering is needed during lowering.
+"""
+
+from typing import Tuple
+
+import jaxlie  # noqa: F401 - validates jaxlie is installed
+
+from ..expr import Expr, to_expr
+
+
+class SE3Exp(Expr):
+    """Exponential map from se(3) twist to SE(3) transformation matrix.
+
+    Maps a 6D twist vector to a 4×4 homogeneous transformation matrix.
+    Uses jaxlie for numerically robust implementation with proper handling
+    of small angles and translations.
+
+    The twist ξ = [v; ω] follows the convention:
+
+    - v: 3D linear velocity component
+    - ω: 3D angular velocity component
+
+    This is the key operation for Product of Exponentials (PoE) forward
+    kinematics in robotic manipulators.
+
+    Attributes:
+        twist: 6D twist vector [v; ω] with shape (6,)
+
+    Example:
+        Product of Exponentials forward kinematics::
+
+            import openscvx as ox
+            import numpy as np
+
+            # Screw axis for revolute joint about z-axis at origin
+            screw_axis = np.array([0, 0, 0, 0, 0, 1])  # [v; ω]
+            theta = ox.State("theta", shape=(1,))
+
+            # Joint transformation
+            T = ox.lie.SE3Exp(ox.Constant(screw_axis) * theta)  # 4×4 matrix
+
+            # Chain multiple joints
+            T_01 = ox.lie.SE3Exp(screw1 * q1)
+            T_12 = ox.lie.SE3Exp(screw2 * q2)
+            T_02 = T_01 @ T_12
+
+        Extract position from transformation::
+
+            T_ee = forward_kinematics(joint_angles)
+            p_ee = T_ee[:3, 3]  # End-effector position
+
+    Note:
+        The twist convention [v; ω] matches jaxlie's SE3 tangent
+        parameterization, so no reordering is performed.
+
+    See Also:
+        - SE3Log: Inverse operation (transformation matrix to twist)
+        - SO3Exp: Rotation-only exponential map
+        - AdjointDual: For dynamics computations with twists
+    """
+
+    def __init__(self, twist):
+        """Initialize SE3 exponential map.
+
+        Args:
+            twist: 6D twist vector [v; ω] with shape (6,)
+        """
+        self.twist = to_expr(twist)
+
+    def children(self):
+        return [self.twist]
+
+    def canonicalize(self) -> "Expr":
+        twist = self.twist.canonicalize()
+        return SE3Exp(twist)
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Check that input is a 6D vector and return output shape.
+
+        Returns:
+            tuple: Shape (4, 4) for the homogeneous transformation matrix
+
+        Raises:
+            ValueError: If twist does not have shape (6,)
+        """
+        twist_shape = self.twist.check_shape()
+        if twist_shape != (6,):
+            raise ValueError(f"SE3Exp expects twist with shape (6,), got {twist_shape}")
+        return (4, 4)
+
+    def __repr__(self):
+        return f"SE3Exp({self.twist!r})"
+
+
+class SE3Log(Expr):
+    """Logarithm map from SE(3) transformation matrix to se(3) twist.
+
+    Maps a 4×4 homogeneous transformation matrix to a 6D twist vector.
+    Uses jaxlie for numerically robust implementation.
+
+    The output twist ξ = [v; ω] follows the convention:
+
+    - v: 3D linear component
+    - ω: 3D angular component (rotation vector)
+
+    This is useful for computing error metrics between poses in optimization.
+
+    Attributes:
+        transform: 4×4 homogeneous transformation matrix with shape (4, 4)
+
+    Example:
+        Compute pose error for trajectory optimization::
+
+            import openscvx as ox
+
+            T_current = forward_kinematics(q)
+            T_target = ox.Parameter("T_target", shape=(4, 4), value=goal_pose)
+
+            # Relative transformation
+            T_error = ox.linalg.inv(T_target) @ T_current
+
+            # Convert to twist for error metric
+            twist_error = ox.lie.SE3Log(T_error)
+            pose_cost = ox.linalg.Norm(twist_error) ** 2
+
+    See Also:
+        - SE3Exp: Inverse operation (twist to transformation matrix)
+        - SO3Log: Rotation-only logarithm map
+    """
+
+    def __init__(self, transform):
+        """Initialize SE3 logarithm map.
+
+        Args:
+            transform: 4×4 homogeneous transformation matrix with shape (4, 4)
+        """
+        self.transform = to_expr(transform)
+
+    def children(self):
+        return [self.transform]
+
+    def canonicalize(self) -> "Expr":
+        transform = self.transform.canonicalize()
+        return SE3Log(transform)
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Check that input is a 4×4 matrix and return output shape.
+
+        Returns:
+            tuple: Shape (6,) for the twist vector
+
+        Raises:
+            ValueError: If transform does not have shape (4, 4)
+        """
+        transform_shape = self.transform.check_shape()
+        if transform_shape != (4, 4):
+            raise ValueError(f"SE3Log expects transform with shape (4, 4), got {transform_shape}")
+        return (6,)
+
+    def __repr__(self):
+        return f"SE3Log({self.transform!r})"