openscvx 0.3.2.dev170__py3-none-any.whl

openscvx/symbolic/expr/lie/so3.py

@@ -0,0 +1,138 @@
+"""SO(3) Lie group operations for rotation matrices.
+
+This module provides exponential and logarithm maps for the SO(3) rotation
+group, enabling axis-angle to rotation matrix conversions and vice versa.
+
+Requires jaxlie: pip install openscvx[lie]
+"""
+
+from typing import Tuple
+
+import jaxlie  # noqa: F401 - validates jaxlie is installed
+
+from ..expr import Expr, to_expr
+
+
+class SO3Exp(Expr):
+    """Exponential map from so(3) to SO(3) rotation matrix.
+
+    Maps a 3D rotation vector (axis-angle representation) to a 3×3 rotation
+    matrix using the Rodrigues formula. Uses jaxlie for numerically robust
+    implementation with proper handling of small angles.
+
+    The rotation vector ω has direction equal to the rotation axis and
+    magnitude equal to the rotation angle in radians.
+
+    Attributes:
+        omega: 3D rotation vector with shape (3,)
+
+    Example:
+        Create a rotation about the z-axis::
+
+            import openscvx as ox
+            import numpy as np
+
+            # 90 degree rotation about z
+            omega = ox.Constant(np.array([0, 0, np.pi/2]))
+            R = ox.lie.SO3Exp(omega)  # 3×3 rotation matrix
+
+        Parameterized rotation for optimization::
+
+            theta = ox.State("theta", shape=(1,))
+            axis = ox.Constant(np.array([0, 0, 1]))  # z-axis
+            R = ox.lie.SO3Exp(axis * theta)
+
+    See Also:
+        - SO3Log: Inverse operation (rotation matrix to rotation vector)
+        - SE3Exp: Full rigid body transformation including translation
+    """
+
+    def __init__(self, omega):
+        """Initialize SO3 exponential map.
+
+        Args:
+            omega: 3D rotation vector (axis × angle) with shape (3,)
+        """
+        self.omega = to_expr(omega)
+
+    def children(self):
+        return [self.omega]
+
+    def canonicalize(self) -> "Expr":
+        omega = self.omega.canonicalize()
+        return SO3Exp(omega)
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Check that input is a 3D vector and return output shape.
+
+        Returns:
+            tuple: Shape (3, 3) for the rotation matrix
+
+        Raises:
+            ValueError: If omega does not have shape (3,)
+        """
+        omega_shape = self.omega.check_shape()
+        if omega_shape != (3,):
+            raise ValueError(f"SO3Exp expects omega with shape (3,), got {omega_shape}")
+        return (3, 3)
+
+    def __repr__(self):
+        return f"SO3Exp({self.omega!r})"
+
+
+class SO3Log(Expr):
+    """Logarithm map from SO(3) rotation matrix to so(3) rotation vector.
+
+    Maps a 3×3 rotation matrix to a 3D rotation vector (axis-angle
+    representation). Uses jaxlie for numerically robust implementation.
+
+    The output rotation vector ω has direction equal to the rotation axis
+    and magnitude equal to the rotation angle in radians.
+
+    Attributes:
+        rotation: 3×3 rotation matrix with shape (3, 3)
+
+    Example:
+        Extract rotation vector from a rotation matrix::
+
+            import openscvx as ox
+
+            R = ox.State("R", shape=(3, 3))  # Rotation matrix state
+            omega = ox.lie.SO3Log(R)  # 3D rotation vector
+
+    See Also:
+        - SO3Exp: Inverse operation (rotation vector to rotation matrix)
+        - SE3Log: Full rigid body transformation logarithm
+    """
+
+    def __init__(self, rotation):
+        """Initialize SO3 logarithm map.
+
+        Args:
+            rotation: 3×3 rotation matrix with shape (3, 3)
+        """
+        self.rotation = to_expr(rotation)
+
+    def children(self):
+        return [self.rotation]
+
+    def canonicalize(self) -> "Expr":
+        rotation = self.rotation.canonicalize()
+        return SO3Log(rotation)
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Check that input is a 3×3 matrix and return output shape.
+
+        Returns:
+            tuple: Shape (3,) for the rotation vector
+
+        Raises:
+            ValueError: If rotation does not have shape (3, 3)
+        """
+        rotation_shape = self.rotation.check_shape()
+        if rotation_shape != (3, 3):
+            raise ValueError(f"SO3Log expects rotation with shape (3, 3), got {rotation_shape}")
+        return (3,)
+
+    def __repr__(self):
+        return f"SO3Log({self.rotation!r})"

openscvx/symbolic/expr/linalg.py

@@ -0,0 +1,279 @@
+"""Linear algebra operations for symbolic expressions.
+
+This module provides essential linear algebra operations for matrix and vector
+manipulation in optimization problems. Operations follow NumPy/JAX conventions
+for shapes and broadcasting behavior.
+
+Key Operations:
+    - **Matrix Operations:**
+        - `Transpose` - Matrix/tensor transposition (swaps last two dimensions)
+        - `Diag` - Construct diagonal matrix from vector
+    - **Reductions:**
+        - `Sum` - Sum all elements of an array (reduces to scalar)
+        - `Norm` - Euclidean (L2) norm and other norms of vectors/matrices
+
+Note:
+    For array manipulation operations like stacking and concatenation, see the
+    `array` module.
+
+Example:
+    Matrix transposition and diagonal matrices::
+
+        import openscvx as ox
+        import numpy as np
+
+        # Transpose a matrix
+        A = ox.State("A", shape=(3, 4))
+        A_T = A.T  # Result shape (4, 3)
+
+        # Create a diagonal matrix
+        v = ox.State("v", shape=(5,))
+        D = ox.Diag(v)  # Result shape (5, 5)
+
+    Reduction operations::
+
+        x = ox.State("x", shape=(3, 4))
+
+        # Sum all elements
+        total = ox.Sum(x)  # Result is scalar
+
+        # Compute norm
+        magnitude = ox.Norm(x)  # Result is scalar
+
+    Computing kinetic energy with norms::
+
+        v = ox.State("v", shape=(3,))  # Velocity vector
+        m = 10.0  # Mass
+        kinetic_energy = 0.5 * m * ox.Norm(v)**2
+"""
+
+import hashlib
+from typing import Tuple
+
+from .expr import Expr, to_expr
+
+
+class Transpose(Expr):
+    """Matrix transpose operation for symbolic expressions.
+
+    Transposes the last two dimensions of an expression. For matrices, this swaps
+    rows and columns. For higher-dimensional arrays, it swaps the last two axes.
+    Scalars and vectors are unchanged by transposition.
+
+    The canonicalization includes an optimization that eliminates double transposes:
+    (A.T).T simplifies to A.
+
+    Attributes:
+        operand: Expression to transpose
+
+    Example:
+        Define Tranpose expressions:
+
+            A = Variable("A", shape=(3, 4))
+            A_T = Transpose(A)  # or A.T, result shape (4, 3)
+            v = Variable("v", shape=(5,))
+            v_T = Transpose(v)  # result shape (5,) - vectors unchanged
+    """
+
+    def __init__(self, operand):
+        """Initialize a transpose operation.
+
+        Args:
+            operand: Expression to transpose
+        """
+        self.operand = to_expr(operand)
+
+    def children(self):
+        return [self.operand]
+
+    def canonicalize(self) -> "Expr":
+        """Canonicalize the operand with double transpose optimization."""
+        operand = self.operand.canonicalize()
+
+        # Double transpose optimization: (A.T).T = A
+        if isinstance(operand, Transpose):
+            return operand.operand
+
+        return Transpose(operand)
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Matrix transpose operation swaps the last two dimensions."""
+        operand_shape = self.operand.check_shape()
+
+        if len(operand_shape) == 0:
+            # Scalar transpose is the scalar itself
+            return ()
+        elif len(operand_shape) == 1:
+            # Vector transpose is the vector itself (row vector remains row vector)
+            return operand_shape
+        elif len(operand_shape) == 2:
+            # Matrix transpose: (m,n) -> (n,m)
+            return (operand_shape[1], operand_shape[0])
+        else:
+            # Higher-dimensional array: transpose last two dimensions
+            # (..., m, n) -> (..., n, m)
+            return operand_shape[:-2] + (operand_shape[-1], operand_shape[-2])
+
+    def __repr__(self):
+        return f"({self.operand!r}).T"
+
+
+class Diag(Expr):
+    """Diagonal matrix construction from a vector.
+
+    Creates a square diagonal matrix from a 1D vector. The vector elements become
+    the diagonal entries, with all off-diagonal entries set to zero. This is
+    analogous to numpy.diag() or jax.numpy.diag().
+
+    Note:
+        Currently only supports creating diagonal matrices from vectors.
+        Extracting diagonals from matrices is not yet implemented.
+
+    Attributes:
+        operand: 1D vector expression to place on the diagonal
+
+    Example:
+        Define a Diag:
+
+            v = Variable("v", shape=(3,))
+            D = Diag(v)  # Creates a (3, 3) diagonal matrix
+    """
+
+    def __init__(self, operand):
+        """Initialize a diagonal matrix operation.
+
+        Args:
+            operand: 1D vector expression to place on the diagonal
+        """
+        self.operand = to_expr(operand)
+
+    def children(self):
+        return [self.operand]
+
+    def canonicalize(self) -> "Expr":
+        operand = self.operand.canonicalize()
+        return Diag(operand)
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Diag converts a vector (n,) to a diagonal matrix (n,n)."""
+        operand_shape = self.operand.check_shape()
+        if len(operand_shape) != 1:
+            raise ValueError(f"Diag expects a 1D vector, got shape {operand_shape}")
+        n = operand_shape[0]
+        return (n, n)
+
+    def __repr__(self):
+        return f"diag({self.operand!r})"
+
+
+class Sum(Expr):
+    """Sum reduction operation for symbolic expressions.
+
+    Sums all elements of an expression, reducing it to a scalar. This is a
+    reduction operation that collapses all dimensions.
+
+    Attributes:
+        operand: Expression whose elements will be summed
+
+    Example:
+        Define a Sum expression::
+
+            x = ox.State("x", shape=(3, 4))
+            total = Sum(x)  # Creates Sum(x), result shape ()
+    """
+
+    def __init__(self, operand):
+        """Initialize a sum reduction operation.
+
+        Args:
+            operand: Expression to sum over all elements
+        """
+        self.operand = to_expr(operand)
+
+    def children(self):
+        return [self.operand]
+
+    def canonicalize(self) -> "Expr":
+        """Canonicalize sum: canonicalize the operand.
+
+        Returns:
+            Expr: Canonical form of the sum expression
+        """
+        operand = self.operand.canonicalize()
+        return Sum(operand)
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Sum reduces any shape to a scalar."""
+        # Validate that the operand has a valid shape
+        self.operand.check_shape()
+        # Sum always produces a scalar regardless of input shape
+        return ()
+
+    def __repr__(self):
+        return f"sum({self.operand!r})"
+
+
+class Norm(Expr):
+    """Norm operation for symbolic expressions (reduction to scalar).
+
+    Computes the norm of an expression according to the specified order parameter.
+    This is a reduction operation that always produces a scalar result regardless
+    of the input shape. Supports various norm types following NumPy/SciPy conventions.
+
+    Attributes:
+        operand: Expression to compute norm of
+        ord: Norm order specification (default: "fro" for Frobenius norm)
+            - "fro": Frobenius norm (default)
+            - "inf": Infinity norm
+            - 1: L1 norm (sum of absolute values)
+            - 2: L2 norm (Euclidean norm)
+            - Other values as supported by the backend
+
+    Example:
+        Define Norms:
+
+            x = Variable("x", shape=(3,))
+            euclidean_norm = Norm(x, ord=2)  # L2 norm, result is scalar
+            A = Variable("A", shape=(3, 4))
+            frobenius_norm = Norm(A)  # Frobenius norm, result is scalar
+    """
+
+    def __init__(self, operand, ord="fro"):
+        """Initialize a norm operation.
+
+        Args:
+            operand: Expression to compute norm of
+            ord: Norm order specification (default: "fro")
+        """
+        self.operand = to_expr(operand)
+        self.ord = ord  # Can be "fro", "inf", 1, 2, etc.
+
+    def children(self):
+        return [self.operand]
+
+    def canonicalize(self) -> "Expr":
+        """Canonicalize the operand but preserve the ord parameter."""
+        canon_operand = self.operand.canonicalize()
+        return Norm(canon_operand, ord=self.ord)
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Norm reduces any shape to a scalar."""
+        # Validate that the operand has a valid shape
+        self.operand.check_shape()
+        # Norm always produces a scalar regardless of input shape
+        return ()
+
+    def _hash_into(self, hasher: "hashlib._Hash") -> None:
+        """Hash Norm including its ord parameter.
+
+        Args:
+            hasher: A hashlib hash object to update
+        """
+        hasher.update(b"Norm")
+        # Hash the ord parameter
+        hasher.update(repr(self.ord).encode())
+        # Hash the operand
+        self.operand._hash_into(hasher)
+
+    def __repr__(self):
+        return f"norm({self.operand!r}, ord={self.ord!r})"