openscvx 0.3.2.dev170__py3-none-any.whl

openscvx/symbolic/expr/array.py

@@ -0,0 +1,632 @@
+"""Array manipulation operations for symbolic expressions.
+
+This module provides operations for indexing, slicing, concatenating, and stacking
+symbolic expressions. These are structural operations that manipulate array shapes
+and combine or extract array elements, as opposed to mathematical transformations.
+
+Key Operations:
+
+- **Indexing and Slicing:**
+    - `Index` - NumPy-style indexing and slicing to extract subarrays
+
+- **Concatenation:**
+    - `Concat` - Concatenate expressions along the first dimension (axis 0)
+
+- **Stacking:**
+    - `Stack` - Stack expressions along a new first dimension
+    - `Hstack` - Horizontal stacking (along columns for 2D arrays)
+    - `Vstack` - Vertical stacking (along rows for 2D arrays)
+
+- **Block Matrix Construction:**
+    - `Block` - Assemble block matrices from nested arrays (like numpy.block)
+
+All operations follow NumPy conventions for shapes and indexing behavior, enabling
+familiar array manipulation patterns in symbolic optimization problems.
+
+Example:
+    Indexing and slicing arrays::
+
+        import openscvx as ox
+
+        x = ox.State("x", shape=(10,))
+        first_half = x[0:5]      # Slice: Index(x, slice(0, 5))
+        element = x[3]           # Single element: Index(x, 3)
+
+        A = ox.State("A", shape=(5, 4))
+        row = A[2, :]            # Extract row
+        col = A[:, 1]            # Extract column
+
+    Concatenating expressions::
+
+        from openscvx.symbolic.expr.array import Concat
+
+        x = ox.State("x", shape=(3,))
+        y = ox.State("y", shape=(4,))
+        combined = Concat(x, y)  # Result shape (7,)
+
+    Stacking to build matrices::
+
+        from openscvx.symbolic.expr.array import Stack, Hstack, Vstack
+
+        # Stack vectors into a matrix
+        v1 = ox.State("v1", shape=(3,))
+        v2 = ox.State("v2", shape=(3,))
+        v3 = ox.State("v3", shape=(3,))
+        matrix = Stack([v1, v2, v3])  # Result shape (3, 3)
+
+        # Horizontal stacking (concatenate along columns)
+        A = ox.State("A", shape=(3, 4))
+        B = ox.State("B", shape=(3, 2))
+        wide = Hstack([A, B])    # Result shape (3, 6)
+
+        # Vertical stacking (concatenate along rows)
+        C = ox.State("C", shape=(2, 4))
+        tall = Vstack([A, C])    # Result shape (5, 4)
+
+    Building rotation matrices with Block (recommended)::
+
+        import openscvx as ox
+        from openscvx.symbolic.expr.array import Block
+
+        theta = ox.Variable("theta", shape=(1,))
+        R = Block([
+            [ox.Cos(theta), -ox.Sin(theta)],
+            [ox.Sin(theta),  ox.Cos(theta)]
+        ])  # 2D rotation matrix, shape (2, 2)
+
+    Building rotation matrices with stacking (alternative)::
+
+        import openscvx as ox
+        from openscvx.symbolic.expr.array import Stack, Hstack
+
+        theta = ox.Variable("theta", shape=(1,))
+        R = Stack([
+            Hstack([ox.Cos(theta), -ox.Sin(theta)]),
+            Hstack([ox.Sin(theta), ox.Cos(theta)])
+        ])  # 2D rotation matrix, shape (2, 2)
+"""
+
+import hashlib
+from typing import Tuple, Union
+
+import numpy as np
+
+from .expr import Expr, to_expr
+
+
+class Index(Expr):
+    """Indexing and slicing operation for symbolic expressions.
+
+    Represents indexing or slicing of an expression using NumPy-style indexing.
+    Can be created using square bracket notation on Expr objects.
+
+    Attributes:
+        base: Expression to index into
+        index: Index specification (int, slice, or tuple of indices/slices)
+
+    Example:
+        Define an Index expression:
+
+            x = ox.State("x", shape=(10,))
+            y = x[0:5]  # Creates Index(x, slice(0, 5))
+            z = x[3]    # Creates Index(x, 3)
+    """
+
+    def __init__(self, base: Expr, index: Union[int, slice, tuple]):
+        """Initialize an indexing operation.
+
+        Args:
+            base: Expression to index into
+            index: NumPy-style index (int, slice, or tuple of indices/slices)
+        """
+        self.base = base
+        self.index = index
+
+    def children(self):
+        return [self.base]
+
+    def canonicalize(self) -> "Expr":
+        """Canonicalize index by canonicalizing the base expression.
+
+        Returns:
+            Expr: Canonical form of the indexing expression
+        """
+        base = self.base.canonicalize()
+        return Index(base, self.index)
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Compute the shape after indexing."""
+        base_shape = self.base.check_shape()
+        dummy = np.zeros(base_shape)
+        try:
+            result = dummy[self.index]
+        except Exception as e:
+            raise ValueError(f"Bad index {self.index} for shape {base_shape}") from e
+        return result.shape
+
+    def _hash_into(self, hasher: "hashlib._Hash") -> None:
+        """Hash Index including its index specification.
+
+        Args:
+            hasher: A hashlib hash object to update
+        """
+        hasher.update(b"Index")
+        # Hash the index specification (convert to string for generality)
+        hasher.update(repr(self.index).encode())
+        # Hash the base expression
+        self.base._hash_into(hasher)
+
+    def __repr__(self):
+        return f"{self.base!r}[{self.index!r}]"
+
+
+class Concat(Expr):
+    """Concatenation operation for symbolic expressions.
+
+    Concatenates a sequence of expressions along their first dimension. All inputs
+    must have the same rank and matching dimensions except for the first dimension.
+
+    Attributes:
+        exprs: Tuple of expressions to concatenate
+
+    Example:
+        Define a Concat expression:
+
+            x = ox.State("x", shape=(3,))
+            y = ox.State("y", shape=(4,))
+            z = Concat(x, y)  # Creates Concat(x, y), result shape (7,)
+    """
+
+    def __init__(self, *exprs: Expr):
+        """Initialize a concatenation operation.
+
+        Args:
+            *exprs: Expressions to concatenate along the first dimension
+        """
+        # wrap raw values as Constant if needed
+        self.exprs = [to_expr(e) for e in exprs]
+
+    def children(self):
+        return list(self.exprs)
+
+    def canonicalize(self) -> "Expr":
+        """Canonicalize concatenation by canonicalizing all operands.
+
+        Returns:
+            Expr: Canonical form of the concatenation expression
+        """
+        exprs = [e.canonicalize() for e in self.exprs]
+        return Concat(*exprs)
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Check concatenation shape compatibility and return result shape."""
+        shapes = [e.check_shape() for e in self.exprs]
+        shapes = [(1,) if len(s) == 0 else s for s in shapes]
+        rank = len(shapes[0])
+        if any(len(s) != rank for s in shapes):
+            raise ValueError(f"Concat rank mismatch: {shapes}")
+        if any(s[1:] != shapes[0][1:] for s in shapes[1:]):
+            raise ValueError(f"Concat non-0 dims differ: {shapes}")
+        return (sum(s[0] for s in shapes),) + shapes[0][1:]
+
+    def __repr__(self):
+        inner = ", ".join(repr(e) for e in self.exprs)
+        return f"Concat({inner})"
+
+
+class Stack(Expr):
+    """Stack expressions vertically to create a higher-dimensional array.
+
+    Stacks a list of expressions along a new first dimension. All input expressions
+    must have the same shape. The result has shape (num_rows, *row_shape).
+
+    This is similar to numpy.array([row1, row2, ...]) or jax.numpy.stack(rows, axis=0).
+
+    Attributes:
+        rows: List of expressions to stack, each representing a "row"
+
+    Example:
+        Leverage stack to combine expressions:
+
+            x = Variable("x", shape=(3,))
+            y = Variable("y", shape=(3,))
+            z = Variable("z", shape=(3,))
+            stacked = Stack([x, y, z])  # Creates shape (3, 3)
+            # Equivalent to: [[x[0], x[1], x[2]],
+            #                 [y[0], y[1], y[2]],
+            #                 [z[0], z[1], z[2]]]
+    """
+
+    def __init__(self, rows):
+        """Initialize a stack operation.
+
+        Args:
+            rows: List of expressions to stack along a new first dimension.
+                  All expressions must have the same shape.
+        """
+        # rows should be a list of expressions representing each row
+        self.rows = [to_expr(row) for row in rows]
+
+    def children(self):
+        return self.rows
+
+    def canonicalize(self) -> "Expr":
+        rows = [row.canonicalize() for row in self.rows]
+        return Stack(rows)
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Stack creates a 2D matrix from 1D rows."""
+        if not self.rows:
+            raise ValueError("Stack requires at least one row")
+
+        # All rows should have the same shape
+        row_shapes = [row.check_shape() for row in self.rows]
+
+        # Verify all rows have the same shape
+        first_shape = row_shapes[0]
+        for i, shape in enumerate(row_shapes[1:], 1):
+            if shape != first_shape:
+                raise ValueError(
+                    f"Stack row {i} has shape {shape}, but row 0 has shape {first_shape}"
+                )
+
+        # Result shape is (num_rows, *row_shape)
+        return (len(self.rows),) + first_shape
+
+    def __repr__(self):
+        rows_repr = ", ".join(repr(row) for row in self.rows)
+        return f"Stack([{rows_repr}])"
+
+
+class Hstack(Expr):
+    """Horizontal stacking operation for symbolic expressions.
+
+    Concatenates expressions horizontally (along columns for 2D arrays).
+    This is analogous to numpy.hstack() or jax.numpy.hstack().
+
+    Behavior depends on input dimensionality:
+    - 1D arrays: Concatenates along axis 0 (making a longer vector)
+    - 2D arrays: Concatenates along axis 1 (columns), rows must match
+    - Higher-D: Concatenates along axis 1, all other dimensions must match
+
+    Attributes:
+        arrays: List of expressions to stack horizontally
+
+    Example:
+        1D case: concatenate vectors:
+
+            x = Variable("x", shape=(3,))
+            y = Variable("y", shape=(2,))
+            h = Hstack([x, y])  # Result shape (5,)
+
+        2D case: concatenate matrices horizontally:
+
+            A = Variable("A", shape=(3, 4))
+            B = Variable("B", shape=(3, 2))
+            C = Hstack([A, B])  # Result shape (3, 6)
+    """
+
+    def __init__(self, arrays):
+        """Initialize a horizontal stack operation.
+
+        Args:
+            arrays: List of expressions to concatenate horizontally
+        """
+        self.arrays = [to_expr(arr) for arr in arrays]
+
+    def children(self):
+        return self.arrays
+
+    def canonicalize(self) -> "Expr":
+        arrays = [arr.canonicalize() for arr in self.arrays]
+        return Hstack(arrays)
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Horizontal stack concatenates arrays along the second axis (columns)."""
+        if not self.arrays:
+            raise ValueError("Hstack requires at least one array")
+
+        array_shapes = [arr.check_shape() for arr in self.arrays]
+
+        # All arrays must have the same number of dimensions
+        first_ndim = len(array_shapes[0])
+        for i, shape in enumerate(array_shapes[1:], 1):
+            if len(shape) != first_ndim:
+                raise ValueError(
+                    f"Hstack array {i} has {len(shape)} dimensions, but array 0 has {first_ndim}"
+                )
+
+        # For 1D arrays, hstack concatenates along axis 0
+        if first_ndim == 1:
+            total_length = sum(shape[0] for shape in array_shapes)
+            return (total_length,)
+
+        # For 2D+ arrays, all dimensions except the second must match
+        first_shape = array_shapes[0]
+        for i, shape in enumerate(array_shapes[1:], 1):
+            if shape[0] != first_shape[0]:
+                raise ValueError(
+                    f"Hstack array {i} has {shape[0]} rows, but array 0 has {first_shape[0]} rows"
+                )
+            if shape[2:] != first_shape[2:]:
+                raise ValueError(
+                    f"Hstack array {i} has trailing dimensions {shape[2:]}, "
+                    f"but array 0 has {first_shape[2:]}"
+                )
+
+        # Result shape: concatenate along axis 1 (columns)
+        total_cols = sum(shape[1] for shape in array_shapes)
+        return (first_shape[0], total_cols) + first_shape[2:]
+
+    def __repr__(self):
+        arrays_repr = ", ".join(repr(arr) for arr in self.arrays)
+        return f"Hstack([{arrays_repr}])"
+
+
+class Vstack(Expr):
+    """Vertical stacking operation for symbolic expressions.
+
+    Concatenates expressions vertically (along rows for 2D arrays).
+    This is analogous to numpy.vstack() or jax.numpy.vstack().
+
+    All input expressions must have the same number of dimensions, and all
+    dimensions except the first must match. The result concatenates along
+    axis 0 (rows).
+
+    Attributes:
+        arrays: List of expressions to stack vertically
+
+    Example:
+        Stack vectors to create a matrix:
+
+            x = Variable("x", shape=(3,))
+            y = Variable("y", shape=(3,))
+            v = Vstack([x, y])  # Result shape (2, 3)
+
+        Stack matrices vertically:
+
+            A = Variable("A", shape=(3, 4))
+            B = Variable("B", shape=(2, 4))
+            C = Vstack([A, B])  # Result shape (5, 4)
+    """
+
+    def __init__(self, arrays):
+        """Initialize a vertical stack operation.
+
+        Args:
+            arrays: List of expressions to concatenate vertically.
+                    All must have matching dimensions except the first.
+        """
+        self.arrays = [to_expr(arr) for arr in arrays]
+
+    def children(self):
+        return self.arrays
+
+    def canonicalize(self) -> "Expr":
+        arrays = [arr.canonicalize() for arr in self.arrays]
+        return Vstack(arrays)
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Vertical stack concatenates arrays along the first axis (rows)."""
+        if not self.arrays:
+            raise ValueError("Vstack requires at least one array")
+
+        array_shapes = [arr.check_shape() for arr in self.arrays]
+
+        # All arrays must have the same number of dimensions
+        first_ndim = len(array_shapes[0])
+        for i, shape in enumerate(array_shapes[1:], 1):
+            if len(shape) != first_ndim:
+                raise ValueError(
+                    f"Vstack array {i} has {len(shape)} dimensions, but array 0 has {first_ndim}"
+                )
+
+        # All dimensions except the first must match
+        first_shape = array_shapes[0]
+        for i, shape in enumerate(array_shapes[1:], 1):
+            if shape[1:] != first_shape[1:]:
+                raise ValueError(
+                    f"Vstack array {i} has trailing dimensions {shape[1:]}, "
+                    f"but array 0 has {first_shape[1:]}"
+                )
+
+        # Result shape: concatenate along axis 0 (rows)
+        total_rows = sum(shape[0] for shape in array_shapes)
+        return (total_rows,) + first_shape[1:]
+
+    def __repr__(self):
+        arrays_repr = ", ".join(repr(arr) for arr in self.arrays)
+        return f"Vstack([{arrays_repr}])"
+
+
+class Block(Expr):
+    """Block matrix/tensor construction from nested arrays of expressions.
+
+    Assembles a block matrix (or N-D tensor) from a nested list of expressions,
+    analogous to numpy.block(). Each inner list represents a row of blocks, and
+    blocks within the same row are concatenated horizontally, while rows are
+    stacked vertically.
+
+    This provides a convenient way to construct matrices from sub-expressions
+    without manually nesting Stack/Hstack/Vstack operations.
+
+    Attributes:
+        blocks: Nested list of expressions forming the block structure (each
+            expression can be a scalar, 1D, 2D, or N-D tensor)
+
+    Example:
+        Build a 2D rotation matrix::
+
+            import openscvx as ox
+            from openscvx.symbolic.expr.array import Block
+
+            theta = ox.Variable("theta", shape=(1,))
+            R = Block([
+                [ox.Cos(theta), -ox.Sin(theta)],
+                [ox.Sin(theta),  ox.Cos(theta)]
+            ])  # Result shape (2, 2)
+
+        Build a block diagonal matrix::
+
+            A = ox.State("A", shape=(2, 2))
+            B = ox.State("B", shape=(3, 3))
+            zeros_23 = ox.Constant(np.zeros((2, 3)))
+            zeros_32 = ox.Constant(np.zeros((3, 2)))
+            block_diag = Block([
+                [A, zeros_23],
+                [zeros_32, B]
+            ])  # Result shape (5, 5)
+
+        Build from scalars and expressions::
+
+            x = ox.State("x", shape=(1,))
+            y = ox.State("y", shape=(1,))
+            # Scalars are automatically promoted to 1D arrays
+            M = Block([
+                [x, 0],
+                [0, y]
+            ])  # Result shape (2, 2)
+
+    Note:
+        - All blocks in the same row must have the same height (first dimension)
+        - All blocks in the same column must have the same width (second dimension)
+        - For N-D tensors (3D+), all trailing dimensions must match across all blocks
+        - Scalar values and raw Python lists are automatically wrapped via to_expr()
+        - 1D arrays are treated as row vectors when determining block dimensions
+        - N-D tensors are supported for JAX lowering; CVXPy only supports 2D blocks
+    """
+
+    def __init__(self, blocks):
+        """Initialize a block matrix construction.
+
+        Args:
+            blocks: A nested list of expressions. Can be either:
+                    - 2D: [[row1_blocks], [row2_blocks], ...] for multiple rows
+                    - 1D: [block1, block2, ...] for a single row (auto-promoted to [[...]])
+                    Raw values (numbers, lists, numpy arrays) are automatically
+                    converted to Constant expressions.
+
+        Raises:
+            ValueError: If blocks is empty
+        """
+        if not blocks:
+            raise ValueError("Block requires at least one row")
+
+        # Auto-promote 1D list to 2D (matching numpy.block behavior)
+        # e.g., Block([a, b]) -> Block([[a, b]])
+        if not isinstance(blocks[0], (list, tuple)):
+            blocks = [blocks]
+
+        # Convert all blocks to expressions
+        self.blocks = [[to_expr(block) for block in row] for row in blocks]
+
+        # Validate consistent row lengths
+        row_lengths = [len(row) for row in self.blocks]
+        if len(set(row_lengths)) > 1:
+            raise ValueError(
+                f"All rows must have the same number of blocks. Got row lengths: {row_lengths}"
+            )
+
+    def children(self):
+        """Return all block expressions in row-major order."""
+        return [block for row in self.blocks for block in row]
+
+    def canonicalize(self) -> "Expr":
+        """Canonicalize by recursively canonicalizing all blocks.
+
+        If the block contains only a single element ([[a]]), returns the
+        canonicalized element directly to simplify the expression tree.
+        """
+        canonical_blocks = [[block.canonicalize() for block in row] for row in self.blocks]
+
+        # Unwrap single-element blocks
+        if len(canonical_blocks) == 1 and len(canonical_blocks[0]) == 1:
+            return canonical_blocks[0][0]
+
+        return Block(canonical_blocks)
+
+    def check_shape(self) -> Tuple[int, ...]:
+        """Validate block dimensions and compute output shape.
+
+        For 2D blocks, returns (total_rows, total_cols). For N-D blocks,
+        returns the shape after assembling blocks along the first two axes,
+        with trailing dimensions preserved.
+
+        Returns:
+            Tuple representing the assembled block array shape
+
+        Raises:
+            ValueError: If block dimensions are incompatible
+        """
+        n_block_rows = len(self.blocks)
+        n_block_cols = len(self.blocks[0])
+
+        # Get shapes of all blocks
+        block_shapes = [[block.check_shape() for block in row] for row in self.blocks]
+
+        # Determine the maximum dimensionality across all blocks
+        max_ndim = max(len(shape) for row in block_shapes for shape in row)
+        max_ndim = max(max_ndim, 2)  # At least 2D for block assembly
+
+        # Normalize shapes: pad to max_ndim by prepending 1s
+        # Scalars () -> (1, 1, ...), 1D (n,) -> (1, n, ...), etc.
+        def normalize_shape(shape):
+            if len(shape) == 0:
+                return (1,) * max_ndim
+            elif len(shape) < max_ndim:
+                # Prepend 1s to match max_ndim
+                return (1,) * (max_ndim - len(shape)) + shape
+            else:
+                return shape
+
+        normalized_shapes = [[normalize_shape(shape) for shape in row] for row in block_shapes]
+
+        # Validate trailing dimensions (dims 2+) match across ALL blocks
+        if max_ndim > 2:
+            trailing_shape = normalized_shapes[0][0][2:]
+            for i, row_shapes in enumerate(normalized_shapes):
+                for j, shape in enumerate(row_shapes):
+                    if shape[2:] != trailing_shape:
+                        raise ValueError(
+                            f"Block[{i}][{j}] has trailing dimensions {shape[2:]}, "
+                            f"but Block[0][0] has {trailing_shape}. "
+                            f"All blocks must have matching dimensions beyond the first two."
+                        )
+
+        # Compute row heights (first dimension of each row must match)
+        row_heights = []
+        for i, row_shapes in enumerate(normalized_shapes):
+            heights = [s[0] for s in row_shapes]
+            if len(set(heights)) > 1:
+                raise ValueError(
+                    f"Block row {i} has inconsistent heights: {heights}. "
+                    f"All blocks in a row must have the same height."
+                )
+            row_heights.append(heights[0])
+
+        # Compute column widths (second dimension of each column must match)
+        col_widths = []
+        for j in range(n_block_cols):
+            widths = [normalized_shapes[i][j][1] for i in range(n_block_rows)]
+            if len(set(widths)) > 1:
+                raise ValueError(
+                    f"Block column {j} has inconsistent widths: {widths}. "
+                    f"All blocks in a column must have the same width."
+                )
+            col_widths.append(widths[0])
+
+        total_rows = sum(row_heights)
+        total_cols = sum(col_widths)
+
+        # Return shape with trailing dimensions if present
+        if max_ndim > 2:
+            return (total_rows, total_cols) + normalized_shapes[0][0][2:]
+        return (total_rows, total_cols)
+
+    def __repr__(self):
+        rows_repr = []
+        for row in self.blocks:
+            blocks_repr = ", ".join(repr(block) for block in row)
+            rows_repr.append(f"[{blocks_repr}]")
+        inner = ", ".join(rows_repr)
+        return f"Block([{inner}])"