openscvx 0.3.2.dev170__py3-none-any.whl

openscvx/symbolic/expr/spatial.py

@@ -0,0 +1,209 @@
+"""Spatial and 6-DOF utility operations for trajectory optimization.
+
+This module provides efficient symbolic expression nodes for common 6-DOF (six degree of
+freedom) operations used in aerospace and robotics applications. These operations directly
+map to optimized JAX implementations for high-performance evaluation.
+"""
+
+from typing import Tuple
+
+from .expr import Expr, to_expr
+
+
+class QDCM(Expr):
+    """Quaternion to Direction Cosine Matrix (DCM) conversion.
+
+    Converts a unit quaternion representation to a 3x3 direction cosine matrix
+    (also known as a rotation matrix). This operation is commonly used in 6-DOF
+    spacecraft dynamics, aircraft simulation, and robotics applications.
+
+    The quaternion is expected in scalar-last format: [qx, qy, qz, qw] where
+    qw is the scalar component. The resulting DCM can be used to transform vectors
+    from one reference frame to another.
+
+    Attributes:
+        q: Quaternion expression with shape (4,)
+
+    Example:
+        Use the QDCM to rotate a vector:
+
+            import openscvx as ox
+            q = ox.State("q", shape=(4,))
+            dcm = ox.QDCM(q)  # Creates rotation matrix, shape (3, 3)
+            v_body = ox.Variable("v_body", shape=(3,))
+            v_inertial = dcm @ v_body
+
+    Note:
+        The input quaternion does not need to be normalized; the implementation
+        automatically handles normalization during evaluation.
+    """
+
+    def __init__(self, q):
+        """Initialize a quaternion to DCM conversion.
+
+        Args:
+            q: Quaternion expression with shape (4,) in [qx, qy, qz, qw] format
+        """
+        self.q = to_expr(q)
+
+    def children(self):
+        return [self.q]
+
+    def canonicalize(self) -> "Expr":
+        q = self.q.canonicalize()
+        return QDCM(q)
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Check that input is a quaternion and return DCM shape.
+
+        Returns:
+            tuple: Shape (3, 3) for the resulting direction cosine matrix
+
+        Raises:
+            ValueError: If quaternion does not have shape (4,)
+        """
+        q_shape = self.q.check_shape()
+        if q_shape != (4,):
+            raise ValueError(f"QDCM expects quaternion with shape (4,), got {q_shape}")
+        return (3, 3)
+
+    def __repr__(self):
+        return f"qdcm({self.q!r})"
+
+
+class SSMP(Expr):
+    """Angular rate to 4x4 skew-symmetric matrix for quaternion dynamics.
+
+    Constructs the 4x4 skew-symmetric matrix Ω(ω) used in quaternion kinematic
+    differential equations. This matrix relates angular velocity to the time
+    derivative of the quaternion:
+
+        q̇ = (1/2) * Ω(ω) @ q
+
+    The resulting matrix has the form:
+        ⎡  0   ωz  -ωy   ωx ⎤
+        ⎢-ωz    0   ωx   ωy ⎥
+        ⎢ ωy  -ωx    0   ωz ⎥
+        ⎣-ωx  -ωy  -ωz    0 ⎦
+
+    This is particularly useful for formulating quaternion-based attitude
+    dynamics in spacecraft and aircraft trajectory optimization problems.
+
+    Attributes:
+        w: Angular velocity vector expression with shape (3,)
+
+    Example:
+        Use the SSMP to compute the quaternion derivative:
+
+            import openscvx as ox
+            omega = ox.Control("omega", shape=(3,))
+            q = ox.State("q", shape=(4,))
+            # Quaternion kinematic equation
+            q_dot = 0.5 * ox.SSMP(omega) @ q
+
+    See Also:
+        SSM: 3x3 skew-symmetric matrix for cross product operations
+    """
+
+    def __init__(self, w):
+        """Initialize an angular velocity to skew-symmetric matrix conversion.
+
+        Args:
+            w: Angular velocity vector expression with shape (3,) in [ωx, ωy, ωz] format
+        """
+        self.w = to_expr(w)
+
+    def children(self):
+        return [self.w]
+
+    def canonicalize(self) -> "Expr":
+        w = self.w.canonicalize()
+        return SSMP(w)
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Check that input is a 3D angular velocity and return matrix shape.
+
+        Returns:
+            tuple: Shape (4, 4) for the resulting skew-symmetric matrix
+
+        Raises:
+            ValueError: If angular velocity does not have shape (3,)
+        """
+        w_shape = self.w.check_shape()
+        if w_shape != (3,):
+            raise ValueError(f"SSMP expects angular velocity with shape (3,), got {w_shape}")
+        return (4, 4)
+
+    def __repr__(self):
+        return f"ssmp({self.w!r})"
+
+
+class SSM(Expr):
+    """Angular rate vector to 3x3 skew-symmetric matrix (cross product matrix).
+
+    Constructs the 3x3 skew-symmetric matrix [ω]x that represents the cross
+    product operation. For any 3D vector v, the cross product ω x v can be
+    computed as the matrix-vector product [ω]x @ v.
+
+    The resulting matrix has the form:
+        ⎡  0  -ωz   ωy ⎤
+        ⎢ ωz    0  -ωx ⎥
+        ⎣-ωy   ωx    0 ⎦
+
+    This operation is widely used in:
+    - Rigid body dynamics (angular momentum calculations)
+    - DCM time derivatives: Ṙ = [ω]x @ R
+    - Velocity kinematics in robotics
+    - Coriolis and centrifugal acceleration terms
+
+    Attributes:
+        w: Angular velocity or 3D vector expression with shape (3,)
+
+    Example:
+        Use the SSM to compute the rotation matrix derivative:
+
+            import openscvx as ox
+            omega = ox.Control("omega", shape=(3,))
+            R = ox.State("R", shape=(3, 3))  # Direction cosine matrix
+            # DCM time derivative
+            R_dot = ox.SSM(omega) @ R
+
+    Note:
+        The skew-symmetric property ensures that [ω]xᵀ = -[ω]x, which is
+        important for preserving orthogonality in DCM propagation.
+
+    See Also:
+        SSMP: 4x4 skew-symmetric matrix for quaternion dynamics
+    """
+
+    def __init__(self, w):
+        """Initialize a vector to skew-symmetric matrix conversion.
+
+        Args:
+            w: 3D vector expression with shape (3,) in [ωx, ωy, ωz] format
+        """
+        self.w = to_expr(w)
+
+    def children(self):
+        return [self.w]
+
+    def canonicalize(self) -> "Expr":
+        w = self.w.canonicalize()
+        return SSM(w)
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Check that input is a 3D vector and return matrix shape.
+
+        Returns:
+            tuple: Shape (3, 3) for the resulting skew-symmetric matrix
+
+        Raises:
+            ValueError: If input vector does not have shape (3,)
+        """
+        w_shape = self.w.check_shape()
+        if w_shape != (3,):
+            raise ValueError(f"SSM expects angular velocity with shape (3,), got {w_shape}")
+        return (3, 3)
+
+    def __repr__(self):
+        return f"ssm({self.w!r})"