openscvx 0.3.2.dev170__py3-none-any.whl

openscvx/symbolic/expr/expr.py

@@ -0,0 +1,720 @@
+"""Core symbolic expression system for trajectory optimization.
+
+This module provides the foundation for openscvx's symbolic expression framework,
+implementing an Abstract Syntax Tree (AST) representation for mathematical expressions
+used in optimization problems. The expression system enables:
+
+- Declarative problem specification: Write optimization problems using familiar
+    mathematical notation with operator overloading (+, -, *, /, @, **, etc.)
+- Automatic differentiation: Expressions are automatically differentiated during
+    compilation to solver-specific formats
+- Shape checking: Static validation of tensor dimensions before optimization
+- Canonicalization: Algebraic simplification for more efficient compilation
+- Multiple backends: Expressions can be compiled to CVXPy, JAX, or custom solvers
+
+Architecture:
+    The expression system is built around an AST where each node is an `Expr` subclass:
+
+    - Leaf nodes: `Parameter`, `Variable`, `State`, `Control` - symbolic values
+    - Arithmetic operations: `Add`, `Sub`, `Mul`, `Div`, `MatMul`, `Power`, `Neg`
+    - Array operations: `Index`, `Concat`, `Stack`, `Hstack`, `Vstack`
+    - Linear algebra: `Transpose`, `Diag`, `Sum`, `Norm`
+    - Constraints: `Equality`, `Inequality`
+    - Functions: `Sin`, `Cos`, `Exp`, `Log`, `Sqrt`, etc.
+
+    Each expression node implements:
+
+    - `children()`: Returns child expressions in the AST
+    - `canonicalize()`: Returns a simplified/normalized version
+    - `check_shape()`: Validates and returns the output shape
+
+Example:
+    Creating symbolic variables and expressions::
+
+        import openscvx as ox
+
+        # Define symbolic variables
+        x = ox.State("x", shape=(3,))
+        A = ox.Parameter("A", shape=(3, 3), value=np.eye(3))
+
+        # Build expressions using natural syntax
+        expr = A @ x + 5
+        constraint = ox.Norm(x) <= 1.0
+
+        # Expressions form an AST
+        print(expr.pretty())  # Visualize the tree structure
+
+    Shape checking with automatic validation::
+
+        x = ox.State("x", shape=(3,))
+        y = ox.State("y", shape=(4,))
+
+        # This will raise ValueError during shape checking
+        try:
+            expr = x + y  # Shapes (3,) and (4,) not broadcastable
+            expr.check_shape()
+        except ValueError as e:
+            print(f"Shape error: {e}")
+
+    Algebraic canonicalization::
+
+        x = ox.State("x", shape=(3,))
+        expr = x + 0 + (1 * x)
+        canonical = expr.canonicalize()  # Simplifies to: x + x
+"""
+
+import hashlib
+import struct
+from typing import Callable, Tuple, Union
+
+import numpy as np
+
+
+class Expr:
+    """Base class for symbolic expressions in optimization problems.
+
+    Expr is the foundation of the symbolic expression system in openscvx. It represents
+    nodes in an abstract syntax tree (AST) for mathematical expressions. Expressions
+    support:
+
+    - Arithmetic operations: +, -, *, /, @, **
+    - Comparison operations: ==, <=, >=
+    - Indexing and slicing: []
+    - Transposition: .T property
+    - Shape checking and validation
+    - Canonicalization (algebraic simplification)
+
+    All Expr subclasses implement a tree structure where each node can have child
+    expressions accessed via the children() method.
+
+    Attributes:
+        __array_priority__: Priority for operations with numpy arrays (set to 1000)
+
+    Note:
+        When used in operations with numpy arrays, Expr objects take precedence,
+        allowing symbolic expressions to wrap numeric values automatically.
+    """
+
+    # Give Expr objects higher priority than numpy arrays in operations
+    __array_priority__ = 1000
+
+    def __le__(self, other):
+        from .constraint import Inequality
+
+        return Inequality(self, to_expr(other))
+
+    def __ge__(self, other):
+        from .constraint import Inequality
+
+        return Inequality(to_expr(other), self)
+
+    def __eq__(self, other):
+        from .constraint import Equality
+
+        return Equality(self, to_expr(other))
+
+    def __add__(self, other):
+        from .arithmetic import Add
+
+        return Add(self, to_expr(other))
+
+    def __radd__(self, other):
+        from .arithmetic import Add
+
+        return Add(to_expr(other), self)
+
+    def __sub__(self, other):
+        from .arithmetic import Sub
+
+        return Sub(self, to_expr(other))
+
+    def __rsub__(self, other):
+        # e.g. 5 - a  ⇒ Sub(Constant(5), a)
+        from .arithmetic import Sub
+
+        return Sub(to_expr(other), self)
+
+    def __truediv__(self, other):
+        from .arithmetic import Div
+
+        return Div(self, to_expr(other))
+
+    def __rtruediv__(self, other):
+        # e.g. 10 / a
+        from .arithmetic import Div
+
+        return Div(to_expr(other), self)
+
+    def __mul__(self, other):
+        from .arithmetic import Mul
+
+        return Mul(self, to_expr(other))
+
+    def __rmul__(self, other):
+        from .arithmetic import Mul
+
+        return Mul(to_expr(other), self)
+
+    def __matmul__(self, other):
+        from .arithmetic import MatMul
+
+        return MatMul(self, to_expr(other))
+
+    def __rmatmul__(self, other):
+        from .arithmetic import MatMul
+
+        return MatMul(to_expr(other), self)
+
+    def __rle__(self, other):
+        # other <= self  =>  Inequality(other, self)
+        from .constraint import Inequality
+
+        return Inequality(to_expr(other), self)
+
+    def __rge__(self, other):
+        # other >= self  =>  Inequality(self, other)
+        from .constraint import Inequality
+
+        return Inequality(self, to_expr(other))
+
+    def __req__(self, other):
+        # other == self  =>  Equality(other, self)
+        from .constraint import Equality
+
+        return Equality(to_expr(other), self)
+
+    def __neg__(self):
+        from .arithmetic import Neg
+
+        return Neg(self)
+
+    def __pow__(self, other):
+        from .arithmetic import Power
+
+        return Power(self, to_expr(other))
+
+    def __rpow__(self, other):
+        from .arithmetic import Power
+
+        return Power(to_expr(other), self)
+
+    def __getitem__(self, idx):
+        from .array import Index
+
+        return Index(self, idx)
+
+    @property
+    def T(self):
+        """Transpose property for matrix expressions.
+
+        Returns:
+            Transpose: A Transpose expression wrapping this expression
+
+        Example:
+            Create a transpose:
+
+                A = ox.State("A", shape=(3, 4))
+                A_T = A.T  # Creates Transpose(A), result shape (4, 3)
+        """
+        from .linalg import Transpose
+
+        return Transpose(self)
+
+    def at(self, k: int) -> "NodeReference":
+        """Reference this expression at a specific trajectory node.
+
+        This method enables inter-node constraints where you can reference
+        the value of an expression at different time steps. Common patterns
+        include rate limits and multi-step dependencies.
+
+        Args:
+            k: Absolute node index (integer) in the trajectory.
+               Can be positive (0, 1, 2, ...) or negative (-1 for last node).
+
+        Returns:
+            NodeReference: An expression representing this expression at node k
+
+        Example:
+            Rate limit constraint (applied across trajectory using a loop):
+
+                position = State("pos", shape=(3,))
+
+                # Create rate limit for each node
+                constraints = [
+                    (ox.linalg.Norm(position.at(k) - position.at(k-1)) <= 0.1).at([k])
+                    for k in range(1, N)
+                ]
+
+            Multi-step dependency:
+
+                state = State("x", shape=(1,))
+
+                # Fibonacci-like recurrence
+                constraints = [
+                    (state.at(k) == state.at(k-1) + state.at(k-2)).at([k])
+                    for k in range(2, N)
+                ]
+
+        Performance Note:
+            Cross-node constraints use dense Jacobian storage which can be memory-intensive
+            for large N (>100 nodes). See LoweredCrossNodeConstraint documentation for
+            details on memory usage and future sparse Jacobian support.
+        """
+        return NodeReference(self, k)
+
+    def children(self):
+        """Return the child expressions of this node.
+
+        Returns:
+            list: List of child Expr objects. Empty list for leaf nodes.
+        """
+        return []
+
+    def canonicalize(self) -> "Expr":
+        """
+        Return a canonical (simplified) form of this expression.
+
+        Canonicalization performs algebraic simplifications such as:
+        - Constant folding (e.g., 2 + 3 → 5)
+        - Identity elimination (e.g., x + 0 → x, x * 1 → x)
+        - Flattening nested operations (e.g., Add(Add(a, b), c) → Add(a, b, c))
+        - Algebraic rewrites (e.g., constraints to standard form)
+
+        Returns:
+            Expr: A canonical version of this expression
+
+        Raises:
+            NotImplementedError: If canonicalization is not implemented for this node type
+        """
+        raise NotImplementedError(f"canonicalize() not implemented for {self.__class__.__name__}")
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """
+        Compute and validate the shape of this expression.
+
+        This method:
+        1. Recursively checks shapes of all child expressions
+        2. Validates that operations are shape-compatible (e.g., broadcasting rules)
+        3. Returns the output shape of this expression
+
+        For example:
+        - A Parameter with shape (3, 4) returns (3, 4)
+        - MatMul of (3, 4) @ (4, 5) returns (3, 5)
+        - Sum of any shape returns () (scalar)
+        - Add broadcasts shapes like NumPy
+
+        Returns:
+            tuple: The shape of this expression as a tuple of integers.
+                   Empty tuple () represents a scalar.
+
+        Raises:
+            NotImplementedError: If shape checking is not implemented for this node type
+            ValueError: If the expression has invalid shapes (e.g., incompatible dimensions)
+        """
+        raise NotImplementedError(f"check_shape() not implemented for {self.__class__.__name__}")
+
+    def pretty(self, indent=0):
+        """Generate a pretty-printed string representation of the expression tree.
+
+        Creates an indented, hierarchical view of the expression tree structure,
+        useful for debugging and visualization.
+
+        Args:
+            indent: Current indentation level (default: 0)
+
+        Returns:
+            str: Multi-line string representation of the expression tree
+
+        Example:
+            Pretty print an expression:
+
+                expr = (x + y) * z
+                print(expr.pretty())
+                # Mul
+                #   Add
+                #     State
+                #     State
+                #   State
+        """
+        pad = "  " * indent
+        pad = "  " * indent
+        lines = [f"{pad}{self.__class__.__name__}"]
+        for child in self.children():
+            lines.append(child.pretty(indent + 1))
+        return "\n".join(lines)
+
+    def _hash_into(self, hasher: "hashlib._Hash") -> None:
+        """Contribute this expression's structural identity to a hash.
+
+        This method is used to compute a structural hash of the expression tree
+        that is name-invariant (same structure = same hash regardless of variable names).
+
+        The default implementation hashes the class name and recursively hashes all
+        children. Subclasses with additional attributes (like Norm.ord, Index.index)
+        should override this to include those attributes.
+
+        Args:
+            hasher: A hashlib hash object to update
+        """
+        # Hash the class name to distinguish different node types
+        hasher.update(self.__class__.__name__.encode())
+        # Recursively hash all children
+        for child in self.children():
+            child._hash_into(hasher)
+
+    def structural_hash(self) -> bytes:
+        """Compute a structural hash of this expression.
+
+        Returns a hash that depends only on the mathematical structure of the
+        expression, not on variable names. Two expressions that are structurally
+        equivalent (same operations, same variable positions) will have the same hash.
+
+        Returns:
+            bytes: SHA-256 digest of the expression structure
+        """
+        hasher = hashlib.sha256()
+        self._hash_into(hasher)
+        return hasher.digest()
+
+
+class Leaf(Expr):
+    """
+    Base class for leaf nodes (terminal expressions) in the symbolic expression tree.
+
+    Leaf nodes represent named symbolic variables that don't have child expressions.
+    This includes Parameters, Variables, States, and Controls.
+
+    Attributes:
+        name (str): Name identifier for the leaf node
+        _shape (tuple): Shape of the leaf node
+    """
+
+    def __init__(self, name: str, shape: tuple = ()):
+        """Initialize a Leaf node.
+
+        Args:
+            name (str): Name identifier for the leaf node
+            shape (tuple): Shape of the leaf node
+        """
+        super().__init__()
+        self.name = name
+        self._shape = shape
+
+    @property
+    def shape(self):
+        """Get the shape of the leaf node.
+
+        Returns:
+            tuple: Shape of the leaf node
+        """
+        return self._shape
+
+    def children(self):
+        """Leaf nodes have no children.
+
+        Returns:
+            list: Empty list since leaf nodes are terminal
+        """
+        return []
+
+    def canonicalize(self) -> "Expr":
+        """Leaf nodes are already in canonical form.
+
+        Returns:
+            Expr: Returns self since leaf nodes are already canonical
+        """
+        return self
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Return the shape of this leaf node.
+
+        Returns:
+            tuple: The shape of the leaf node
+        """
+        return self._shape
+
+    def _hash_into(self, hasher: "hashlib._Hash") -> None:
+        """Hash leaf node by class name and shape.
+
+        This base implementation hashes the class name and shape. Subclasses
+        like Variable and Parameter override this to add their specific
+        canonical identifiers (_slice for Variables, value for Parameters).
+
+        Args:
+            hasher: A hashlib hash object to update
+        """
+        hasher.update(self.__class__.__name__.encode())
+        hasher.update(str(self._shape).encode())
+
+    def __repr__(self):
+        """String representation of the leaf node.
+
+        Returns:
+            str: A string describing the leaf node
+        """
+        return f"{self.__class__.__name__}('{self.name}', shape={self.shape})"
+
+
+class Parameter(Leaf):
+    """Parameter that can be changed at runtime without recompilation.
+
+    Parameters are symbolic variables with initial values that can be updated
+    through the problem's parameter dictionary. They allow for efficient
+    parameter sweeps without needing to recompile the optimization problem.
+
+    Example:
+        obs_center = ox.Parameter("obs_center", shape=(3,), value=np.array([1.0, 0.0, 0.0]))
+        # Later: problem.parameters["obs_center"] = new_value
+    """
+
+    def __init__(self, name: str, shape: tuple = (), value=None):
+        """Initialize a Parameter node.
+
+        Args:
+            name (str): Name identifier for the parameter
+            shape (tuple): Shape of the parameter (default: scalar)
+            value: Initial value for the parameter (required)
+        """
+        super().__init__(name, shape)
+        if value is None:
+            raise ValueError(f"Parameter '{name}' requires an initial value")
+        self.value = np.asarray(value, dtype=float)
+
+    def _hash_into(self, hasher: "hashlib._Hash") -> None:
+        """Hash Parameter by its shape only (value-invariant).
+
+        Parameters are hashed by shape only, not by value. This allows the same
+        compiled solver to be reused across parameter sweeps - only the structure
+        matters for compilation, not the actual values.
+
+        Args:
+            hasher: A hashlib hash object to update
+        """
+        hasher.update(b"Parameter")
+        hasher.update(str(self._shape).encode())
+
+
+def to_expr(x: Union[Expr, float, int, np.ndarray]) -> Expr:
+    """Convert a value to an Expr if it is not already one.
+
+    This is a convenience function that wraps numeric values and arrays as Constant
+    expressions, while leaving Expr instances unchanged. Used internally by operators
+    to ensure operands are proper Expr objects.
+
+    Args:
+        x: Value to convert - can be an Expr, numeric scalar, or numpy array
+
+    Returns:
+        The input if it's already an Expr, otherwise a Constant wrapping the value
+    """
+    return x if isinstance(x, Expr) else Constant(np.array(x))
+
+
+def traverse(expr: Expr, visit: Callable[[Expr], None]):
+    """Depth-first traversal of an expression tree.
+
+    Visits each node in the expression tree by applying the visit function to the
+    current node, then recursively visiting all children.
+
+    Args:
+        expr: Root expression node to start traversal from
+        visit: Callback function applied to each node during traversal
+    """
+    visit(expr)
+    for child in expr.children():
+        traverse(child, visit)
+
+
+class Constant(Expr):
+    """Constant value expression.
+
+    Represents a constant numeric value in the expression tree. Constants are
+    automatically normalized (squeezed) upon construction to ensure consistency.
+
+    Attributes:
+        value: The numpy array representing the constant value (squeezed)
+
+    Example:
+        Define constants:
+
+            c1 = Constant(5.0)        # Scalar constant
+            c2 = Constant([1, 2, 3])  # Vector constant
+            c3 = to_expr(10)          # Also creates a Constant
+    """
+
+    def __init__(self, value: np.ndarray):
+        """Initialize a constant expression.
+
+        Args:
+            value: Numeric value or numpy array to wrap as a constant.
+                   Will be converted to numpy array and squeezed.
+        """
+        # Normalize immediately upon construction to ensure consistency
+        # This ensures Constant(5.0) and Constant([5.0]) create identical objects
+        if not isinstance(value, np.ndarray):
+            value = np.array(value, dtype=float)
+        self.value = np.squeeze(value)
+
+    def canonicalize(self) -> "Expr":
+        """Constants are already in canonical form.
+
+        Returns:
+            Expr: Returns self since constants are already canonical
+        """
+        return self
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Return the shape of this constant's value.
+
+        Returns:
+            tuple: The shape of the constant's numpy array value
+        """
+        # Verify the invariant: constants should already be squeezed during construction
+        original_shape = self.value.shape
+        squeezed_shape = np.squeeze(self.value).shape
+        if original_shape != squeezed_shape:
+            raise ValueError(
+                f"Constant not properly normalized: has shape {original_shape} "
+                "but should have shape {squeezed_shape}. "
+                "Constants should be squeezed during construction."
+            )
+        return self.value.shape
+
+    def _hash_into(self, hasher: "hashlib._Hash") -> None:
+        """Hash constant by its value.
+
+        Constants are hashed by their actual numeric value, ensuring that
+        expressions with the same constant values produce the same hash.
+
+        Args:
+            hasher: A hashlib hash object to update
+        """
+        hasher.update(b"Constant")
+        hasher.update(str(self.value.shape).encode())
+        hasher.update(self.value.tobytes())
+
+    def __repr__(self):
+        # Show clean representation - always show as Python values, not numpy arrays
+        if self.value.ndim == 0:
+            # Scalar: show as plain number
+            return f"Const({self.value.item()!r})"
+        else:
+            # Array: show as Python list for readability
+            return f"Const({self.value.tolist()!r})"
+
+
+class NodeReference(Expr):
+    """Reference to a variable at a specific trajectory node.
+
+    NodeReference enables inter-node constraints by allowing you to reference
+    the value of a state or control variable at a specific discrete time point
+    (node) in the trajectory. This is essential for expressing temporal relationships
+    such as:
+
+    - Rate limits and smoothness constraints
+    - Multi-step dependencies and recurrence relations
+    - Constraints coupling specific nodes
+
+    Attributes:
+        base: The expression (typically a Leaf like State or Control) being referenced
+        node_idx: Trajectory node index (integer, can be negative for end-indexing)
+
+    Example:
+        Rate limit across trajectory:
+
+            position = State("pos", shape=(3,))
+
+            # Create rate limit constraints for all nodes
+            constraints = [
+                (ox.linalg.Norm(position.at(k) - position.at(k-1)) <= 0.1).at([k])
+                for k in range(1, N)
+            ]
+
+        Multi-step dependency:
+
+            state = State("x", shape=(1,))
+
+            # Fibonacci-like recurrence at each node
+            constraints = [
+                (state.at(k) == state.at(k-1) + state.at(k-2)).at([k])
+                for k in range(2, N)
+            ]
+
+        Coupling specific nodes:
+
+            # Constrain distance between nodes 5 and 10
+            coupling = (position.at(10) - position.at(5) <= threshold).at([10])
+
+    Performance Note:
+        Cross-node constraints use dense Jacobian storage. For details on memory
+        usage and performance implications, see LoweredCrossNodeConstraint documentation.
+
+    Note:
+        NodeReference is typically created via the `.at(k)` method on expressions
+        rather than constructed directly.
+    """
+
+    def __init__(self, base: Expr, node_idx: int):
+        """Initialize a NodeReference.
+
+        Args:
+            base: Expression to reference at a specific node (typically a Leaf)
+            node_idx: Absolute trajectory node index (integer)
+                     Supports negative indexing (e.g., -1 for last node)
+
+        Raises:
+            TypeError: If node_idx is not an integer
+        """
+        if not isinstance(node_idx, int):
+            raise TypeError(f"Node index must be an integer, got {type(node_idx).__name__}")
+
+        self.node_idx = node_idx
+        self.base = base
+
+    def children(self):
+        """Return the base expression as the only child.
+
+        Returns:
+            list: Single-element list containing the base expression
+        """
+        return [self.base]
+
+    def canonicalize(self) -> "Expr":
+        """Canonicalize by canonicalizing the base expression.
+
+        Returns:
+            NodeReference: A new NodeReference with canonicalized base
+        """
+        canon_base = self.base.canonicalize()
+        return NodeReference(canon_base, self.node_idx)
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Return the shape of the base expression.
+
+        NodeReference doesn't change the shape of the underlying expression,
+        it just references it at a specific time point.
+
+        Returns:
+            tuple: The shape of the base expression
+        """
+        return self.base.check_shape()
+
+    def _hash_into(self, hasher: "hashlib._Hash") -> None:
+        """Hash NodeReference including its node index.
+
+        Args:
+            hasher: A hashlib hash object to update
+        """
+        hasher.update(b"NodeReference")
+        # Hash the node index (signed int)
+        hasher.update(struct.pack(">i", self.node_idx))
+        # Hash the base expression
+        self.base._hash_into(hasher)
+
+    def __repr__(self):
+        """String representation of the NodeReference.
+
+        Returns:
+            str: String showing the base expression and node index
+        """
+        return f"{self.base!r}.at({self.node_idx})"